akar_common/
memory_account.rs

1//! Memory accounting classes and attribution.
2//!
3//! The buffer manager and index structures allocate/deallocate against
4//! [`crate::memory::MemoryManager`]. Attribution lets the same manager track
5//! *why* memory is used (which subsystem/class), so the effective spill
6//! threshold can be derived from live usage instead of a fixed constant.
7
8use std::collections::HashMap;
9use std::sync::Mutex;
10use std::sync::atomic::{AtomicU64, Ordering};
11
12/// Accounting class = which subsystem caused the allocation.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum MemoryAccountingClass {
15    /// Buffer-manager page cache frames.
16    BufferPool,
17    /// Vector (HNSW) indexes.
18    VectorIndex,
19    /// Full-text indexes.
20    FtsIndex,
21    /// Graph adjacency structures.
22    Graph,
23    /// Compiled execution plans.
24    CompiledPlan,
25    /// Anything not covered above.
26    Other,
27}
28
29impl MemoryAccountingClass {
30    /// Stable machine-readable label.
31    pub fn as_str(self) -> &'static str {
32        match self {
33            MemoryAccountingClass::BufferPool => "buffer_pool",
34            MemoryAccountingClass::VectorIndex => "vector_index",
35            MemoryAccountingClass::FtsIndex => "fts_index",
36            MemoryAccountingClass::Graph => "graph",
37            MemoryAccountingClass::CompiledPlan => "compiled_plan",
38            MemoryAccountingClass::Other => "other",
39        }
40    }
41
42    /// Stable 64-bit discriminator (for Cheap-to-store accounting).
43    pub fn discriminant(self) -> u8 {
44        match self {
45            MemoryAccountingClass::BufferPool => 0,
46            MemoryAccountingClass::VectorIndex => 1,
47            MemoryAccountingClass::FtsIndex => 2,
48            MemoryAccountingClass::Graph => 3,
49            MemoryAccountingClass::CompiledPlan => 4,
50            MemoryAccountingClass::Other => 5,
51        }
52    }
53}
54
55impl From<u8> for MemoryAccountingClass {
56    fn from(v: u8) -> Self {
57        match v {
58            0 => MemoryAccountingClass::BufferPool,
59            1 => MemoryAccountingClass::VectorIndex,
60            2 => MemoryAccountingClass::FtsIndex,
61            3 => MemoryAccountingClass::Graph,
62            4 => MemoryAccountingClass::CompiledPlan,
63            _ => MemoryAccountingClass::Other,
64        }
65    }
66}
67
68/// Attribution for one allocation: which subsystem/domain and accounting class.
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub struct MemoryAttribution {
71    /// Free-form subsystem/domain name, e.g. `"buffer_pool"` or `"catalog"`.
72    pub domain: &'static str,
73    /// Accounting class for bucketed reporting.
74    pub class: MemoryAccountingClass,
75}
76
77/// Canonical attribution for buffer-pool page frames.
78pub const BUFFER_POOL: MemoryAttribution = MemoryAttribution {
79    domain: "buffer_pool",
80    class: MemoryAccountingClass::BufferPool,
81};
82
83/// Canonical attribution for vector (HNSW/ANN) indexes.
84pub const VECTOR_INDEX: MemoryAttribution = MemoryAttribution {
85    domain: "vector_index",
86    class: MemoryAccountingClass::VectorIndex,
87};
88
89/// Canonical attribution for full-text indexes.
90pub const FTS_INDEX: MemoryAttribution = MemoryAttribution {
91    domain: "fts_index",
92    class: MemoryAccountingClass::FtsIndex,
93};
94
95/// Per-class memory usage bookkeeping.
96///
97/// O(1) allocate/deallocate and cheap snapshots; safe to share via `Mutex`
98/// since these sit behind the [`crate::memory::MemoryManager`] handle.
99#[derive(Debug, Default)]
100pub struct MemoryAccountant {
101    total: AtomicU64,
102    by_class: Mutex<HashMap<MemoryAccountingClass, u64>>,
103}
104
105impl MemoryAccountant {
106    /// Create an empty accountant.
107    pub fn new() -> Self {
108        Self::default()
109    }
110
111    /// Attribute `amount` bytes to `attr`.
112    pub fn allocate(&self, attr: MemoryAttribution, amount: u64) {
113        self.total.fetch_add(amount, Ordering::Relaxed);
114        let mut by_class = self.by_class.lock().unwrap();
115        *by_class.entry(attr.class).or_insert(0) += amount;
116    }
117
118    /// Release `amount` bytes attributed to `attr` (saturates per-class so an
119    /// over-release can never make totals negative).
120    pub fn deallocate(&self, attr: MemoryAttribution, amount: u64) {
121        self.total
122            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |prev| {
123                Some(prev.saturating_sub(amount))
124            })
125            .ok();
126        let mut by_class = self.by_class.lock().unwrap();
127        if let Some(bytes) = by_class.get_mut(&attr.class) {
128            *bytes = bytes.saturating_sub(amount);
129        }
130    }
131
132    /// Total attributed bytes currently live.
133    pub fn total(&self) -> u64 {
134        self.total.load(Ordering::Relaxed)
135    }
136
137    /// Attributed bytes for one class (0 if never allocated).
138    pub fn class_usage(&self, class: MemoryAccountingClass) -> u64 {
139        self.by_class.lock().unwrap().get(&class).copied().unwrap_or(0)
140    }
141
142    /// Snapshot of `(class, bytes)` for reporting, non-zero classes only.
143    pub fn snapshot(&self) -> Vec<(MemoryAccountingClass, u64)> {
144        let mut out: Vec<(MemoryAccountingClass, u64)> = self
145            .by_class
146            .lock()
147            .unwrap()
148            .iter()
149            .map(|(&class, &bytes)| (class, bytes))
150            .filter(|(_, bytes)| *bytes > 0)
151            .collect();
152        out.sort_by_key(|(class, _)| class.discriminant());
153        out
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    #[test]
162    fn test_allocate_deallocate_roundtrip() {
163        let acc = MemoryAccountant::new();
164        acc.allocate(BUFFER_POOL, 4096);
165        acc.allocate(VECTOR_INDEX, 2048);
166        assert_eq!(acc.total(), 6144);
167        assert_eq!(acc.class_usage(MemoryAccountingClass::BufferPool), 4096);
168        acc.deallocate(BUFFER_POOL, 4096);
169        assert_eq!(acc.total(), 2048);
170        assert_eq!(acc.class_usage(MemoryAccountingClass::BufferPool), 0);
171    }
172
173    #[test]
174    fn test_deallocate_never_negative() {
175        let acc = MemoryAccountant::new();
176        acc.allocate(BUFFER_POOL, 100);
177        acc.deallocate(BUFFER_POOL, 10_000);
178        assert_eq!(acc.class_usage(MemoryAccountingClass::BufferPool), 0);
179        assert_eq!(acc.total(), 0);
180    }
181
182    #[test]
183    fn test_snapshot_sorted_non_zero() {
184        let acc = MemoryAccountant::new();
185        acc.allocate(FTS_INDEX, 128);
186        acc.allocate(BUFFER_POOL, 256);
187        let snap = acc.snapshot();
188        assert_eq!(
189            snap,
190            vec![
191                (MemoryAccountingClass::BufferPool, 256),
192                (MemoryAccountingClass::FtsIndex, 128),
193            ]
194        );
195    }
196
197    #[test]
198    fn test_class_discriminants_roundtrip() {
199        for class in [
200            MemoryAccountingClass::BufferPool,
201            MemoryAccountingClass::VectorIndex,
202            MemoryAccountingClass::FtsIndex,
203            MemoryAccountingClass::Graph,
204            MemoryAccountingClass::CompiledPlan,
205            MemoryAccountingClass::Other,
206        ] {
207            assert_eq!(MemoryAccountingClass::from(class.discriminant()), class);
208        }
209    }
210}
akar_common/memory_account.rs

akar_common/
memory_account.rs
