akar_common/
memory.rs

1//! Memory management utilities for the buffer manager.
2//!
3//! Tracks allocated memory and provides backpressure hints.
4
5use crate::memory_account::{MemoryAccountant, MemoryAccountingClass, MemoryAttribution};
6use std::sync::atomic::{AtomicU64, Ordering};
7
8/// A memory tracker for the database instance.
9///
10/// In addition to a flat `total_allocated` counter, the manager routes every
11/// allocation through a [`MemoryAccountant`] so the same budget can be broken
12/// down by subsystem (buffer pool, indexes, graphs) and used to derive an
13/// *effective* spill threshold for the memory governor.
14#[derive(Debug)]
15pub struct MemoryManager {
16    total_allocated: AtomicU64,
17    max_memory: u64,
18    accountant: MemoryAccountant,
19}
20
21impl MemoryManager {
22    pub fn new(max_memory: u64) -> Self {
23        Self {
24            total_allocated: AtomicU64::new(0),
25            max_memory,
26            accountant: MemoryAccountant::new(),
27        }
28    }
29
30    pub fn max_memory(&self) -> u64 {
31        self.max_memory
32    }
33
34    pub fn total_allocated(&self) -> u64 {
35        self.total_allocated.load(Ordering::Relaxed)
36    }
37
38    /// Allocate `amount` bytes attributed to [`MemoryAccountingClass::Other`].
39    ///
40    /// Kept for callers that do not care about attribution; prefer
41    /// [`MemoryManager::allocate_with`].
42    pub fn allocate(&self, amount: u64) {
43        self.allocate_with(
44            MemoryAttribution {
45                domain: "other",
46                class: MemoryAccountingClass::Other,
47            },
48            amount,
49        );
50    }
51
52    /// Allocate `amount` bytes attributed to `attr` (accountable allocation).
53    pub fn allocate_with(&self, attr: MemoryAttribution, amount: u64) {
54        self.total_allocated.fetch_add(amount, Ordering::Relaxed);
55        self.accountant.allocate(attr, amount);
56    }
57
58    /// Release `amount` bytes (unattributed). Prefer
59    /// [`MemoryManager::deallocate_with`] to keep per-class books balanced.
60    pub fn deallocate(&self, amount: u64) {
61        self.deallocate_with(
62            MemoryAttribution {
63                domain: "other",
64                class: MemoryAccountingClass::Other,
65            },
66            amount,
67        );
68    }
69
70    /// Release `amount` bytes previously attributed to `attr`.
71    pub fn deallocate_with(&self, attr: MemoryAttribution, amount: u64) {
72        self.total_allocated.fetch_sub(amount, Ordering::Relaxed);
73        self.accountant.deallocate(attr, amount);
74    }
75
76    pub fn is_under_limit(&self) -> bool {
77        self.total_allocated() <= self.max_memory
78    }
79
80    /// The accountant backing this manager (live totals break down by class).
81    pub fn accountant(&self) -> &MemoryAccountant {
82        &self.accountant
83    }
84
85    // -------------------------------------------------------------------
86    // Memory-governor hooks (effective spill threshold & pressure)
87    // -------------------------------------------------------------------
88
89    /// Headroom left before the configured budget is exhausted.
90    ///
91    /// Intended to be read by the spiller when deciding whether a NodeGroup /
92    /// frame batch crosses the spill threshold: "how much can I still grow?".
93    pub fn effective_spill_threshold(&self) -> u64 {
94        self.max_memory.saturating_sub(self.total_allocated())
95    }
96
97    /// Fraction of the configured budget currently allocated, in `0.0..=1.0`.
98    pub fn memory_pressure(&self) -> f64 {
99        if self.max_memory == 0 {
100            return 0.0;
101        }
102        (self.total_allocated() as f64 / self.max_memory as f64).clamp(0.0, 1.0)
103    }
104
105    /// True when allocated memory consumes at least `ratio` of the budget.
106    pub fn is_under_memory_pressure(&self, ratio: f64) -> bool {
107        self.memory_pressure() >= ratio
108    }
109}
110
111impl Default for MemoryManager {
112    fn default() -> Self {
113        // Default to 80% of available system memory (approximate).
114        Self::new(u64::MAX)
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use crate::memory_account::{BUFFER_POOL, MemoryAccountingClass};
122
123    #[test]
124    fn test_allocate_with_is_accounted() {
125        let mm = MemoryManager::new(1024);
126        mm.allocate_with(BUFFER_POOL, 100);
127        assert_eq!(mm.total_allocated(), 100);
128        assert_eq!(mm.accountant().class_usage(MemoryAccountingClass::BufferPool), 100);
129        mm.deallocate_with(BUFFER_POOL, 100);
130        assert_eq!(mm.total_allocated(), 0);
131    }
132
133    #[test]
134    fn test_plain_allocate_uses_other_class() {
135        let mm = MemoryManager::new(1024);
136        mm.allocate(64);
137        assert_eq!(mm.accountant().class_usage(MemoryAccountingClass::Other), 64);
138    }
139
140    #[test]
141    fn test_effective_spill_threshold() {
142        let mm = MemoryManager::new(1000);
143        assert_eq!(mm.effective_spill_threshold(), 1000);
144        mm.allocate_with(BUFFER_POOL, 300);
145        assert_eq!(mm.effective_spill_threshold(), 700);
146        // Saturates rather than going negative.
147        mm.allocate_with(BUFFER_POOL, 10_000);
148        assert_eq!(mm.effective_spill_threshold(), 0);
149    }
150
151    #[test]
152    fn test_memory_pressure() {
153        let mm = MemoryManager::new(100);
154        assert!((mm.memory_pressure() - 0.0).abs() < 1e-9);
155        assert!(!mm.is_under_memory_pressure(0.8));
156        mm.allocate_with(BUFFER_POOL, 90);
157        assert!((mm.memory_pressure() - 0.9).abs() < 1e-9);
158        assert!(mm.is_under_memory_pressure(0.8));
159    }
160}
akar_common/memory.rs

akar_common/
memory.rs
