nodedb_mem/
scoped_memory.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Governor handle pre-bound to a database, tenant, and engine.
4//!
5//! A call site holding one cannot omit a budget layer. The identity is set
6//! at construction, never per call.
7
8use std::sync::Arc;
9
10use nodedb_types::{DatabaseId, TenantId};
11
12use crate::engine::EngineId;
13use crate::error::Result;
14use crate::governor::MemoryGovernor;
15use crate::reservation_token::ReservationToken;
16
17/// A [`MemoryGovernor`] handle bound to one database, tenant, and engine.
18#[derive(Clone, Debug)]
19pub struct ScopedMemory {
20    governor: Arc<MemoryGovernor>,
21    database_id: DatabaseId,
22    tenant_id: TenantId,
23    engine: EngineId,
24}
25
26impl ScopedMemory {
27    pub fn new(
28        governor: Arc<MemoryGovernor>,
29        database_id: DatabaseId,
30        tenant_id: TenantId,
31        engine: EngineId,
32    ) -> Self {
33        Self {
34            governor,
35            database_id,
36            tenant_id,
37            engine,
38        }
39    }
40
41    /// Reserve `bytes`, checking global, database, tenant, then engine.
42    ///
43    /// # Errors
44    ///
45    /// Names the first exhausted layer.
46    pub fn reserve(&self, bytes: usize) -> Result<ReservationToken> {
47        self.governor
48            .try_reserve(self.database_id, self.tenant_id, self.engine, bytes)
49    }
50
51    /// Account `bytes` already resident. Never fails.
52    ///
53    /// Denying a charge hides allocated memory, it never frees it.
54    pub fn charge(&self, bytes: usize) -> ReservationToken {
55        self.governor
56            .charge(self.database_id, self.tenant_id, self.engine, bytes)
57    }
58
59    pub fn engine(&self) -> EngineId {
60        self.engine
61    }
62
63    pub fn database_id(&self) -> DatabaseId {
64        self.database_id
65    }
66
67    pub fn tenant_id(&self) -> TenantId {
68        self.tenant_id
69    }
70
71    pub fn governor(&self) -> &Arc<MemoryGovernor> {
72        &self.governor
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use std::sync::Arc;
79    use std::sync::atomic::Ordering;
80
81    use nodedb_types::{DatabaseId, TenantId};
82
83    use super::*;
84    use crate::engine_limits::EngineLimits;
85    use crate::error::MemError;
86    use crate::governor::GovernorConfig;
87
88    fn make_governor(limits: &[(EngineId, usize)], ceiling: usize) -> Arc<MemoryGovernor> {
89        let engine_limits: EngineLimits = limits.iter().copied().collect();
90        Arc::new(
91            MemoryGovernor::new(GovernorConfig {
92                global_ceiling: ceiling,
93                engine_limits,
94            })
95            .expect("valid config"),
96        )
97    }
98
99    fn scoped(gov: &Arc<MemoryGovernor>, engine: EngineId) -> ScopedMemory {
100        ScopedMemory::new(
101            Arc::clone(gov),
102            DatabaseId::DEFAULT,
103            TenantId::new(1),
104            engine,
105        )
106    }
107
108    #[test]
109    fn reservation_moves_the_global_counter_symmetrically() {
110        let gov = make_governor(&[(EngineId::Vector, 4096)], 8192);
111        let mem = scoped(&gov, EngineId::Vector);
112
113        {
114            let _token = mem.reserve(1000).expect("within budget");
115            assert_eq!(
116                gov.global_counter.allocated.load(Ordering::Relaxed),
117                1000,
118                "a reservation must increment the global counter it releases on drop"
119            );
120        }
121
122        assert_eq!(
123            gov.global_counter.allocated.load(Ordering::Relaxed),
124            0,
125            "dropping the token must return the global counter to zero"
126        );
127    }
128
129    #[test]
130    fn reserve_over_engine_budget_returns_err_and_charges_nothing() {
131        let gov = make_governor(&[(EngineId::Fts, 512)], 1024);
132        let mem = scoped(&gov, EngineId::Fts);
133
134        let err = mem.reserve(1000).unwrap_err();
135        assert!(
136            matches!(err, MemError::BudgetExhausted { .. }),
137            "expected BudgetExhausted, got {err:?}"
138        );
139        assert_eq!(gov.budget(EngineId::Fts).allocated(), 0);
140    }
141
142    #[test]
143    fn independent_engines_accumulate_and_release_independently() {
144        let gov = make_governor(
145            &[
146                (EngineId::Vector, 4096),
147                (EngineId::Columnar, 4096),
148                (EngineId::Graph, 4096),
149            ],
150            16384,
151        );
152        let vector = scoped(&gov, EngineId::Vector);
153        let columnar = scoped(&gov, EngineId::Columnar);
154        let graph = scoped(&gov, EngineId::Graph);
155
156        let t1 = vector.reserve(1000).unwrap();
157        let t2 = columnar.reserve(2000).unwrap();
158        let t3 = graph.reserve(3000).unwrap();
159
160        assert_eq!(gov.budget(EngineId::Vector).allocated(), 1000);
161        assert_eq!(gov.budget(EngineId::Columnar).allocated(), 2000);
162        assert_eq!(gov.budget(EngineId::Graph).allocated(), 3000);
163
164        drop(t2);
165        assert_eq!(gov.budget(EngineId::Vector).allocated(), 1000);
166        assert_eq!(gov.budget(EngineId::Columnar).allocated(), 0);
167        assert_eq!(gov.budget(EngineId::Graph).allocated(), 3000);
168
169        drop(t1);
170        drop(t3);
171        assert_eq!(gov.budget(EngineId::Vector).allocated(), 0);
172        assert_eq!(gov.budget(EngineId::Graph).allocated(), 0);
173    }
174
175    #[test]
176    fn freed_reservation_can_be_reserved_again() {
177        let gov = make_governor(&[(EngineId::Timeseries, 1024)], 2048);
178        let mem = scoped(&gov, EngineId::Timeseries);
179
180        {
181            let _t = mem.reserve(1024).unwrap();
182            assert!(mem.reserve(1).is_err(), "budget fully consumed");
183        } // token dropped -> budget freed
184
185        let _t2 = mem
186            .reserve(1024)
187            .expect("budget freed by previous token drop");
188    }
189
190    /// `mem::forget` on the returned token prevents release across all
191    /// four layers `try_reserve` charges.
192    #[test]
193    fn mem_forget_does_not_release() {
194        let gov = make_governor(&[(EngineId::Kv, 4096)], 8192);
195        let mem = scoped(&gov, EngineId::Kv);
196
197        let token = mem.reserve(500).unwrap();
198        assert_eq!(gov.budget(EngineId::Kv).allocated(), 500);
199
200        std::mem::forget(token);
201
202        assert_eq!(
203            gov.budget(EngineId::Kv).allocated(),
204            500,
205            "mem::forget intentionally skips drop; bytes remain charged"
206        );
207    }
208
209    #[test]
210    fn reserve_zero_bytes_is_allowed() {
211        let gov = make_governor(&[(EngineId::Query, 1024)], 2048);
212        let mem = scoped(&gov, EngineId::Query);
213
214        let token = mem.reserve(0).expect("zero bytes always fits");
215        assert_eq!(token.size(), 0);
216        drop(token);
217        assert_eq!(gov.budget(EngineId::Query).allocated(), 0);
218    }
219
220    // ── `charge`: infallible accounting for already-resident memory ─────────
221
222    #[test]
223    fn charge_past_the_engine_limit_still_accounts_the_bytes() {
224        let gov = make_governor(&[(EngineId::Timeseries, 100)], 1000);
225        let mem = scoped(&gov, EngineId::Timeseries);
226
227        let token = mem.charge(500);
228        assert_eq!(token.size(), 500);
229        assert_eq!(
230            gov.budget(EngineId::Timeseries).allocated(),
231            500,
232            "a charge past the engine limit must still be visible to the governor"
233        );
234    }
235
236    #[test]
237    fn charge_past_the_global_ceiling_still_accounts_and_clamps_utilization() {
238        let gov = make_governor(&[(EngineId::Timeseries, 50)], 100);
239        let mem = scoped(&gov, EngineId::Timeseries);
240
241        let token = mem.charge(500);
242        assert_eq!(token.size(), 500);
243        assert_eq!(
244            gov.global_utilization_percent(),
245            100,
246            "utilization must clamp at 100%, not wrap, when a charge exceeds the ceiling"
247        );
248    }
249
250    #[test]
251    fn dropping_a_charged_token_releases_every_layer() {
252        let gov = make_governor(&[(EngineId::Kv, 1024)], 2048);
253        let mem = scoped(&gov, EngineId::Kv);
254
255        {
256            let _token = mem.charge(2000);
257            assert_eq!(gov.budget(EngineId::Kv).allocated(), 2000);
258            assert_eq!(gov.global_counter.allocated.load(Ordering::Relaxed), 2000);
259        }
260
261        assert_eq!(
262            gov.budget(EngineId::Kv).allocated(),
263            0,
264            "dropping a charged token must return the engine counter to baseline"
265        );
266        assert_eq!(
267            gov.global_counter.allocated.load(Ordering::Relaxed),
268            0,
269            "dropping a charged token must return the global counter to baseline"
270        );
271    }
272
273    #[test]
274    fn charge_does_not_loosen_reserve() {
275        let gov = make_governor(&[(EngineId::Graph, 100)], 1000);
276        let mem = scoped(&gov, EngineId::Graph);
277
278        let _held = mem.charge(500);
279        let err = mem
280            .reserve(1)
281            .expect_err("reserve must still refuse past the limit after an unrelated charge");
282        assert!(matches!(err, MemError::BudgetExhausted { .. }));
283    }
284}
nodedb_mem/scoped_memory.rs

nodedb_mem/
scoped_memory.rs
