contextdb_core/
memory.rs

1use std::sync::atomic::Ordering;
2use std::sync::atomic::{AtomicBool, AtomicUsize};
3
4/// Budget enforcer for memory-constrained edge devices.
5/// All methods are &self — interior mutability via atomics.
6#[derive(Debug)]
7#[allow(dead_code)]
8pub struct MemoryAccountant {
9    limit: AtomicUsize,
10    used: AtomicUsize,
11    startup_ceiling: AtomicUsize,
12    has_ceiling: AtomicBool,
13}
14
15#[derive(Debug, Clone)]
16pub struct MemoryUsage {
17    pub limit: Option<usize>,
18    pub used: usize,
19    pub available: Option<usize>,
20    pub startup_ceiling: Option<usize>,
21}
22
23impl MemoryAccountant {
24    /// No budget enforcement. All allocations succeed. Default behavior.
25    pub fn no_limit() -> Self {
26        Self {
27            limit: AtomicUsize::new(0),
28            used: AtomicUsize::new(0),
29            startup_ceiling: AtomicUsize::new(0),
30            has_ceiling: AtomicBool::new(false),
31        }
32    }
33
34    /// Set a byte ceiling. Allocations exceeding this fail.
35    pub fn with_budget(bytes: usize) -> Self {
36        Self {
37            limit: AtomicUsize::new(bytes),
38            used: AtomicUsize::new(0),
39            startup_ceiling: AtomicUsize::new(bytes),
40            has_ceiling: AtomicBool::new(true),
41        }
42    }
43
44    /// Attempt to allocate bytes. CAS-based, no TOCTOU.
45    pub fn try_allocate(&self, bytes: usize) -> crate::Result<()> {
46        if bytes == 0 {
47            return Ok(());
48        }
49
50        loop {
51            let used = self.used.load(Ordering::SeqCst);
52            let limit = self.limit.load(Ordering::SeqCst);
53            if limit != 0 {
54                let available = limit.saturating_sub(used);
55                if bytes > available {
56                    if self.limit.load(Ordering::SeqCst) != limit {
57                        continue;
58                    }
59                    return Err(crate::Error::MemoryBudgetExceeded {
60                        subsystem: "memory".to_string(),
61                        operation: "allocate".to_string(),
62                        requested_bytes: bytes,
63                        available_bytes: available,
64                        budget_limit_bytes: limit,
65                        hint:
66                            "Reduce retained data, lower working-set size, or raise MEMORY_LIMIT."
67                                .to_string(),
68                    });
69                }
70            }
71
72            let next = used.saturating_add(bytes);
73            if self
74                .used
75                .compare_exchange(used, next, Ordering::SeqCst, Ordering::SeqCst)
76                .is_ok()
77            {
78                return Ok(());
79            }
80        }
81    }
82
83    /// Return freed bytes to the budget.
84    pub fn release(&self, bytes: usize) {
85        if bytes == 0 {
86            return;
87        }
88
89        let _ = self
90            .used
91            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |used| {
92                Some(used.saturating_sub(bytes))
93            });
94    }
95
96    /// Runtime budget adjustment. None removes limit.
97    /// Returns Err if new limit exceeds startup ceiling.
98    pub fn set_budget(&self, limit: Option<usize>) -> crate::Result<()> {
99        if self.has_ceiling.load(Ordering::SeqCst) {
100            let ceiling = self.startup_ceiling.load(Ordering::SeqCst);
101            match limit {
102                Some(bytes) if bytes > ceiling => {
103                    return Err(crate::Error::Other(format!(
104                        "memory limit {bytes} exceeds startup ceiling {ceiling}"
105                    )));
106                }
107                None => {
108                    return Err(crate::Error::Other(
109                        "cannot remove memory limit when a startup ceiling is set".to_string(),
110                    ));
111                }
112                _ => {}
113            }
114        }
115
116        match limit {
117            Some(bytes) => {
118                self.limit.store(bytes, Ordering::SeqCst);
119            }
120            None => {
121                self.limit.store(0, Ordering::SeqCst);
122            }
123        }
124
125        Ok(())
126    }
127
128    /// Snapshot of current memory state.
129    pub fn usage(&self) -> MemoryUsage {
130        let limit = match self.limit.load(Ordering::SeqCst) {
131            0 => None,
132            bytes => Some(bytes),
133        };
134        let used = self.used.load(Ordering::SeqCst);
135        let startup_ceiling = self
136            .has_ceiling
137            .load(Ordering::SeqCst)
138            .then(|| self.startup_ceiling.load(Ordering::SeqCst));
139        MemoryUsage {
140            limit,
141            used,
142            available: limit.map(|limit| limit.saturating_sub(used)),
143            startup_ceiling,
144        }
145    }
146
147    pub fn try_allocate_for(
148        &self,
149        bytes: usize,
150        subsystem: &str,
151        operation: &str,
152        hint: &str,
153    ) -> crate::Result<()> {
154        self.try_allocate(bytes).map_err(|err| match err {
155            crate::Error::MemoryBudgetExceeded {
156                requested_bytes,
157                budget_limit_bytes,
158                available_bytes,
159                ..
160            } => crate::Error::MemoryBudgetExceeded {
161                subsystem: subsystem.to_string(),
162                operation: operation.to_string(),
163                requested_bytes,
164                budget_limit_bytes,
165                available_bytes,
166                hint: hint.to_string(),
167            },
168            other => other,
169        })
170    }
171}
contextdb_core/memory.rs

contextdb_core/
memory.rs
