saluki_core/accounting/
mod.rs

1//! Process-level memory bounds and limiting for components.
2//!
3//! This module lets components declare their expected memory usage and supports enforcing
4//! process-wide memory limits:
5//!
6//! - **memory bounds**: components declare their _expected_ memory usage (a minimum required
7//!   amount and a firm limit) via the [`MemoryBounds`] trait and [`MemoryBoundsBuilder`],
8//! - **bounds verification**: [`BoundsVerifier`] checks that the combined bounds of all components
9//!   fit within a [`MemoryGrant`],
10//! - **memory limiting**: [`MemoryLimiter`] applies cooperative backpressure as the process
11//!   approaches a configured memory limit,
12//! - **component registry**: [`ComponentRegistry`] ties these together, providing a nestable
13//!   structure of components that can each declare bounds and register a resource group for
14//!   runtime usage tracking.
15//!
16//! Actual (as opposed to expected) memory usage and CPU time are tracked separately via the
17//! resource-tracking primitives in [`saluki_common::resource_tracking`]; the tracking allocator
18//! found there must be installed as the global allocator for runtime usage to be attributed to
19//! registered components.
20
21use std::collections::HashMap;
22
23use serde::Serialize;
24
25mod api;
26pub use self::api::ResourceAPIHandler;
27
28mod grant;
29pub use self::grant::MemoryGrant;
30
31mod limiter;
32pub use self::limiter::MemoryLimiter;
33
34mod registry;
35pub use self::registry::{ComponentRegistry, ComponentRegistryHandle, MemoryBoundsBuilder};
36
37mod verifier;
38pub use self::verifier::{BoundsVerifier, VerifiedBounds, VerifierError};
39
40#[cfg(test)]
41pub(crate) mod test_util;
42
43/// Memory bounds for a component.
44///
45/// Components will naturally allocate memory in many phases, from initialization to normal operation. In some cases,
46/// these allocations can be unbounded, leading to potential memory exhaustion.
47///
48/// When a component has a way to bound its memory usage, it can implement this trait to provide that accounting. A
49/// bounds builder exposes a simple interface for tallying up the memory usage of individual pieces of a component, such
50/// as buffers and buffer pools, containers, and more.
51pub trait MemoryBounds {
52    /// Specifies the minimum and firm memory bounds for this component and its subcomponents.
53    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder);
54}
55
56impl<T> MemoryBounds for &T
57where
58    T: MemoryBounds,
59{
60    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
61        T::specify_bounds(self, builder);
62    }
63}
64
65impl<T> MemoryBounds for Box<T>
66where
67    T: MemoryBounds + ?Sized,
68{
69    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
70        T::specify_bounds(self, builder);
71    }
72}
73
74/// Represents a memory usage expression for a component.
75#[derive(Clone, Debug, Serialize)]
76#[serde(tag = "type")]
77pub enum UsageExpr {
78    /// A config value
79    Config {
80        /// The name
81        name: String,
82        /// The value
83        value: usize,
84    },
85
86    /// A struct size
87    StructSize {
88        /// The value
89        name: String,
90        /// The value
91        value: usize,
92    },
93
94    /// A constant value
95    Constant {
96        /// The name
97        name: String,
98        /// The value
99        value: usize,
100    },
101
102    /// A product of subexpressions
103    Product {
104        /// Values to multiply
105        values: Vec<UsageExpr>,
106    },
107
108    /// A sum of subexpressions
109    Sum {
110        /// Values to add
111        values: Vec<UsageExpr>,
112    },
113}
114
115impl UsageExpr {
116    /// Creates a new usage expression that's a config value.
117    pub fn config(s: impl Into<String>, value: usize) -> Self {
118        Self::Config { name: s.into(), value }
119    }
120
121    /// Creates a new usage expression that's a constant value.
122    pub fn constant(s: impl Into<String>, value: usize) -> Self {
123        Self::Constant { name: s.into(), value }
124    }
125
126    /// Creates a new usage expression that's a struct size.
127    pub fn struct_size<T>(s: impl Into<String>) -> Self {
128        Self::StructSize {
129            name: s.into(),
130            value: std::mem::size_of::<T>(),
131        }
132    }
133
134    /// Creates a new usage expression that's the product of two subexpressions.
135    pub fn product(_s: impl Into<String>, lhs: UsageExpr, rhs: UsageExpr) -> Self {
136        Self::Product { values: vec![lhs, rhs] }
137    }
138
139    /// Creates a new usage expression that's the sum of two subexpressions.
140    pub fn sum(_s: impl Into<String>, lhs: UsageExpr, rhs: UsageExpr) -> Self {
141        Self::Sum { values: vec![lhs, rhs] }
142    }
143
144    fn evaluate(&self) -> usize {
145        match self {
146            Self::Config { value, .. } | Self::StructSize { value, .. } | Self::Constant { value, .. } => *value,
147            Self::Product { values } => values.iter().map(UsageExpr::evaluate).product(),
148            Self::Sum { values } => values.iter().map(UsageExpr::evaluate).sum(),
149        }
150    }
151}
152
153/// Memory bounds for a component.
154#[derive(Clone, Debug, Default)]
155pub struct ComponentBounds {
156    self_minimum_required_bytes: Vec<UsageExpr>,
157    self_firm_limit_bytes: Vec<UsageExpr>,
158    subcomponents: HashMap<String, ComponentBounds>,
159}
160
161impl ComponentBounds {
162    /// Gets the total minimum required bytes for this component and all subcomponents.
163    pub fn total_minimum_required_bytes(&self) -> usize {
164        self.self_minimum_required_bytes
165            .iter()
166            .map(UsageExpr::evaluate)
167            .sum::<usize>()
168            + self
169                .subcomponents
170                .values()
171                .map(|cb| cb.total_minimum_required_bytes())
172                .sum::<usize>()
173    }
174
175    /// Gets the total firm limit bytes for this component and all subcomponents.
176    ///
177    /// The firm limit includes the minimum required bytes.
178    pub fn total_firm_limit_bytes(&self) -> usize {
179        self.self_minimum_required_bytes
180            .iter()
181            .map(UsageExpr::evaluate)
182            .sum::<usize>()
183            + self
184                .self_firm_limit_bytes
185                .iter()
186                .map(UsageExpr::evaluate)
187                .sum::<usize>()
188            + self
189                .subcomponents
190                .values()
191                .map(|cb| cb.total_firm_limit_bytes())
192                .sum::<usize>()
193    }
194
195    /// Returns an iterator of all subcomponents within this component.
196    ///
197    /// Only iterates over direct subcomponents, not the subcomponents of those subcomponents, and so on.
198    pub fn subcomponents(&self) -> impl IntoIterator<Item = (&String, &ComponentBounds)> {
199        self.subcomponents.iter()
200    }
201
202    /// Returns a tree of all bound expressions for this component and its subcomponents as JSON.
203    pub fn to_exprs(&self) -> Vec<serde_json::Value> {
204        let path = vec!["root".to_string()];
205        let mut stack = vec![(path, self)];
206        let mut output = Vec::new();
207
208        while let Some((path, cb)) = stack.pop() {
209            for expr in &cb.self_minimum_required_bytes {
210                output.push(serde_json::json!({
211                    "name": format!("{}.min", path.join(".")),
212                    "expr": expr,
213                }));
214            }
215            for expr in &cb.self_firm_limit_bytes {
216                output.push(serde_json::json!({
217                    "name": format!("{}.firm", path.join(".")),
218                    "expr": expr,
219                }));
220            }
221
222            for (name, subcomponent) in cb.subcomponents() {
223                let mut path = path.clone();
224                path.push(name.clone());
225                stack.push((path, subcomponent));
226            }
227        }
228
229        output
230    }
231}
232
233#[cfg(test)]
234mod tests {
235    use super::UsageExpr;
236
237    #[test]
238    fn leaf_expressions_evaluate_to_their_value() {
239        assert_eq!(UsageExpr::config("cfg", 7).evaluate(), 7);
240        assert_eq!(UsageExpr::constant("const", 11).evaluate(), 11);
241        assert_eq!(
242            UsageExpr::struct_size::<u64>("u64").evaluate(),
243            std::mem::size_of::<u64>()
244        );
245    }
246
247    #[test]
248    fn product_evaluates_to_the_product_of_its_subexpressions() {
249        let expr = UsageExpr::product(
250            "area",
251            UsageExpr::constant("width", 4),
252            UsageExpr::constant("height", 8),
253        );
254        assert_eq!(expr.evaluate(), 32);
255    }
256
257    #[test]
258    fn sum_evaluates_to_the_sum_of_its_subexpressions() {
259        let expr = UsageExpr::sum("total", UsageExpr::constant("a", 4), UsageExpr::constant("b", 8));
260        assert_eq!(expr.evaluate(), 12);
261    }
262
263    #[test]
264    fn products_and_sums_compose_recursively() {
265        // (2 + 3) * 4 = 20
266        let expr = UsageExpr::product(
267            "scaled",
268            UsageExpr::sum("base", UsageExpr::constant("a", 2), UsageExpr::constant("b", 3)),
269            UsageExpr::constant("factor", 4),
270        );
271        assert_eq!(expr.evaluate(), 20);
272    }
273}
saluki_core/accounting/mod.rs

saluki_core/accounting/
mod.rs
