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    /// Evaluates this expression to a byte count.
145    ///
146    /// Every leaf is ultimately operator-controlled, so an expression can describe a size larger
147    /// than `usize` can hold. Arithmetic saturates at [`usize::MAX`] rather than overflowing: a bound
148    /// too large to represent is reported as the largest one that can be, which no memory grant can
149    /// satisfy, so bounds verification rejects it. Wrapping would instead report a small bound and
150    /// let verification pass.
151    fn evaluate(&self) -> usize {
152        match self {
153            Self::Config { value, .. } | Self::StructSize { value, .. } | Self::Constant { value, .. } => *value,
154            Self::Product { values } => values.iter().map(UsageExpr::evaluate).fold(1, usize::saturating_mul),
155            Self::Sum { values } => values.iter().map(UsageExpr::evaluate).fold(0, usize::saturating_add),
156        }
157    }
158}
159
160/// Memory bounds for a component.
161#[derive(Clone, Debug, Default)]
162pub struct ComponentBounds {
163    self_minimum_required_bytes: Vec<UsageExpr>,
164    self_firm_limit_bytes: Vec<UsageExpr>,
165    subcomponents: HashMap<String, ComponentBounds>,
166}
167
168impl ComponentBounds {
169    /// Gets the total minimum required bytes for this component and all subcomponents.
170    pub fn total_minimum_required_bytes(&self) -> usize {
171        self.self_minimum_required_bytes
172            .iter()
173            .map(UsageExpr::evaluate)
174            .chain(self.subcomponents.values().map(|cb| cb.total_minimum_required_bytes()))
175            .fold(0, usize::saturating_add)
176    }
177
178    /// Gets the total firm limit bytes for this component and all subcomponents.
179    ///
180    /// The firm limit includes the minimum required bytes.
181    pub fn total_firm_limit_bytes(&self) -> usize {
182        self.self_minimum_required_bytes
183            .iter()
184            .chain(self.self_firm_limit_bytes.iter())
185            .map(UsageExpr::evaluate)
186            .chain(self.subcomponents.values().map(|cb| cb.total_firm_limit_bytes()))
187            .fold(0, usize::saturating_add)
188    }
189
190    /// Returns an iterator of all subcomponents within this component.
191    ///
192    /// Only iterates over direct subcomponents, not the subcomponents of those subcomponents, and so on.
193    pub fn subcomponents(&self) -> impl IntoIterator<Item = (&String, &ComponentBounds)> {
194        self.subcomponents.iter()
195    }
196
197    /// Returns a tree of all bound expressions for this component and its subcomponents as JSON.
198    pub fn to_exprs(&self) -> Vec<serde_json::Value> {
199        let path = vec!["root".to_string()];
200        let mut stack = vec![(path, self)];
201        let mut output = Vec::new();
202
203        while let Some((path, cb)) = stack.pop() {
204            for expr in &cb.self_minimum_required_bytes {
205                output.push(serde_json::json!({
206                    "name": format!("{}.min", path.join(".")),
207                    "expr": expr,
208                }));
209            }
210            for expr in &cb.self_firm_limit_bytes {
211                output.push(serde_json::json!({
212                    "name": format!("{}.firm", path.join(".")),
213                    "expr": expr,
214                }));
215            }
216
217            for (name, subcomponent) in cb.subcomponents() {
218                let mut path = path.clone();
219                path.push(name.clone());
220                stack.push((path, subcomponent));
221            }
222        }
223
224        output
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use std::collections::HashMap;
231
232    use super::{ComponentBounds, UsageExpr};
233
234    #[test]
235    fn leaf_expressions_evaluate_to_their_value() {
236        assert_eq!(UsageExpr::config("cfg", 7).evaluate(), 7);
237        assert_eq!(UsageExpr::constant("const", 11).evaluate(), 11);
238        assert_eq!(
239            UsageExpr::struct_size::<u64>("u64").evaluate(),
240            std::mem::size_of::<u64>()
241        );
242    }
243
244    #[test]
245    fn product_evaluates_to_the_product_of_its_subexpressions() {
246        let expr = UsageExpr::product(
247            "area",
248            UsageExpr::constant("width", 4),
249            UsageExpr::constant("height", 8),
250        );
251        assert_eq!(expr.evaluate(), 32);
252    }
253
254    #[test]
255    fn sum_evaluates_to_the_sum_of_its_subexpressions() {
256        let expr = UsageExpr::sum("total", UsageExpr::constant("a", 4), UsageExpr::constant("b", 8));
257        assert_eq!(expr.evaluate(), 12);
258    }
259
260    #[test]
261    fn products_and_sums_compose_recursively() {
262        // (2 + 3) * 4 = 20
263        let expr = UsageExpr::product(
264            "scaled",
265            UsageExpr::sum("base", UsageExpr::constant("a", 2), UsageExpr::constant("b", 3)),
266            UsageExpr::constant("factor", 4),
267        );
268        assert_eq!(expr.evaluate(), 20);
269    }
270
271    #[test]
272    fn empty_products_and_sums_keep_their_identities() {
273        // Folding by hand has to reproduce what `Iterator::product` and `Iterator::sum` return for an
274        // empty sequence, or an expression with no subexpressions changes meaning.
275        assert_eq!(UsageExpr::Product { values: Vec::new() }.evaluate(), 1);
276        assert_eq!(UsageExpr::Sum { values: Vec::new() }.evaluate(), 0);
277    }
278
279    #[test]
280    fn a_sum_too_large_to_represent_saturates() {
281        // An operator-supplied byte budget can be large enough that the total does not fit in a
282        // `usize`. Wrapping would report a tiny bound that verification happily accepts.
283        let expr = UsageExpr::sum(
284            "total",
285            UsageExpr::config("queue budget", usize::MAX),
286            UsageExpr::constant("overhead", 4096),
287        );
288
289        assert_eq!(expr.evaluate(), usize::MAX);
290    }
291
292    #[test]
293    fn a_product_too_large_to_represent_saturates() {
294        let expr = UsageExpr::product(
295            "scaled",
296            UsageExpr::config("count", usize::MAX),
297            UsageExpr::constant("size", 2),
298        );
299
300        assert_eq!(expr.evaluate(), usize::MAX);
301    }
302
303    #[test]
304    fn saturation_survives_aggregation_across_subcomponents() {
305        // Saturating inside `evaluate` is not enough on its own: the per-component totals add those
306        // results together, so an already-saturated expression must not overflow one level up.
307        let saturated = ComponentBounds {
308            self_minimum_required_bytes: vec![UsageExpr::config("queue budget", usize::MAX)],
309            self_firm_limit_bytes: vec![UsageExpr::constant("overhead", 4096)],
310            subcomponents: HashMap::new(),
311        };
312        let mut subcomponents = HashMap::new();
313        subcomponents.insert("forwarder".to_string(), saturated);
314
315        let bounds = ComponentBounds {
316            self_minimum_required_bytes: vec![UsageExpr::constant("root min", 1024)],
317            self_firm_limit_bytes: vec![UsageExpr::constant("root firm", 2048)],
318            subcomponents,
319        };
320
321        assert_eq!(bounds.total_minimum_required_bytes(), usize::MAX);
322        assert_eq!(bounds.total_firm_limit_bytes(), usize::MAX);
323    }
324}
saluki_core/accounting/mod.rs

saluki_core/accounting/
mod.rs
