linkrs_storage/engine/
resource_budget.rs

1//! Resource budgets and memory accounting for the storage engine.
2
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::Arc;
5
6use linkrs_core::error::storage::StorageErrorKind;
7use linkrs_core::types::Timestamp;
8use linkrs_core::{StorageError, StorageResult};
9
10/// Memory ownership categories used by the storage engine.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12#[repr(usize)]
13pub enum MemoryCategory {
14    Data = 0,
15    Index = 1,
16    Mvcc = 2,
17    Cache = 3,
18    Background = 4,
19}
20
21impl MemoryCategory {
22    const COUNT: usize = 5;
23
24    /// Return all categories in a stable order for diagnostics.
25    pub const fn all() -> [Self; Self::COUNT] {
26        [
27            Self::Data,
28            Self::Index,
29            Self::Mvcc,
30            Self::Cache,
31            Self::Background,
32        ]
33    }
34
35    /// Return the metric-friendly category name.
36    pub const fn name(self) -> &'static str {
37        match self {
38            Self::Data => "data",
39            Self::Index => "index",
40            Self::Mvcc => "mvcc",
41            Self::Cache => "cache",
42            Self::Background => "background",
43        }
44    }
45}
46
47/// Validated memory limits for one storage instance.
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub struct MemoryBudget {
50    pub max_memory_bytes: u64,
51    pub index_memory_bytes: u64,
52    pub soft_limit_bytes: u64,
53    pub hard_limit_bytes: u64,
54}
55
56impl MemoryBudget {
57    /// Build a budget from a total limit and soft/hard ratios.
58    pub fn new(
59        max_memory_bytes: u64,
60        index_memory_bytes: u64,
61        soft_ratio: f64,
62        hard_ratio: f64,
63    ) -> StorageResult<Self> {
64        if max_memory_bytes == 0 {
65            return Err(StorageError::new(
66                StorageErrorKind::InvalidInput,
67                "max_memory_bytes must be greater than 0",
68            ));
69        }
70        if index_memory_bytes == 0 || index_memory_bytes > max_memory_bytes {
71            return Err(StorageError::new(
72                StorageErrorKind::InvalidInput,
73                "index_memory_bytes must be greater than 0 and no greater than max_memory_bytes",
74            ));
75        }
76        if !soft_ratio.is_finite()
77            || !hard_ratio.is_finite()
78            || !(0.0..=1.0).contains(&soft_ratio)
79            || !(0.0..=1.0).contains(&hard_ratio)
80            || soft_ratio == 0.0
81            || soft_ratio >= hard_ratio
82        {
83            return Err(StorageError::new(
84                StorageErrorKind::InvalidInput,
85                "memory ratios must satisfy 0 < soft_ratio < hard_ratio <= 1",
86            ));
87        }
88
89        let soft_limit_bytes = (max_memory_bytes as f64 * soft_ratio).round() as u64;
90        let hard_limit_bytes = (max_memory_bytes as f64 * hard_ratio).round() as u64;
91        if soft_limit_bytes == 0 || hard_limit_bytes <= soft_limit_bytes {
92            return Err(StorageError::new(
93                StorageErrorKind::InvalidInput,
94                "memory ratios produce invalid byte limits",
95            ));
96        }
97
98        Ok(Self {
99            max_memory_bytes,
100            index_memory_bytes,
101            soft_limit_bytes,
102            hard_limit_bytes,
103        })
104    }
105
106    pub(crate) fn from_validated(
107        max_memory_bytes: u64,
108        index_memory_bytes: u64,
109        soft_ratio: f64,
110        hard_ratio: f64,
111    ) -> Self {
112        Self {
113            max_memory_bytes,
114            index_memory_bytes,
115            soft_limit_bytes: (max_memory_bytes as f64 * soft_ratio).round() as u64,
116            hard_limit_bytes: (max_memory_bytes as f64 * hard_ratio).round() as u64,
117        }
118    }
119
120    fn category_limit(self, category: MemoryCategory) -> u64 {
121        match category {
122            MemoryCategory::Index => self.index_memory_bytes,
123            _ => self.max_memory_bytes,
124        }
125    }
126}
127
128#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
129pub struct MemoryUsage {
130    pub current_bytes: u64,
131    pub peak_bytes: u64,
132}
133
134/// Point-in-time resource information suitable for logs and metrics exports.
135#[derive(Debug, Clone, Copy, PartialEq, Eq)]
136pub struct ResourceSnapshot {
137    pub budget: MemoryBudget,
138    pub categories: [MemoryUsage; MemoryCategory::COUNT],
139    pub total_current_bytes: u64,
140    pub total_peak_bytes: u64,
141    pub soft_limit_events: u64,
142    pub hard_limit_rejections: u64,
143    pub active_snapshots: usize,
144    pub oldest_snapshot_ts: Timestamp,
145    pub tombstone_count: usize,
146    pub tombstone_memory_bytes: u64,
147    pub vertex_shard_count: usize,
148}
149
150impl ResourceSnapshot {
151    pub fn usage(self, category: MemoryCategory) -> MemoryUsage {
152        self.categories[category as usize]
153    }
154
155    pub fn soft_limit_exceeded(self) -> bool {
156        self.total_current_bytes >= self.budget.soft_limit_bytes
157    }
158
159    pub fn hard_limit_exceeded(self) -> bool {
160        self.total_current_bytes >= self.budget.hard_limit_bytes
161    }
162}
163
164/// Thread-safe accounting shared by storage components.
165pub struct MemoryAccounting {
166    budget: MemoryBudget,
167    current: [AtomicU64; MemoryCategory::COUNT],
168    peaks: [AtomicU64; MemoryCategory::COUNT],
169    total_current: AtomicU64,
170    total_peak: AtomicU64,
171    soft_limit_events: AtomicU64,
172    hard_limit_rejections: AtomicU64,
173}
174
175impl std::fmt::Debug for MemoryAccounting {
176    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177        f.debug_struct("MemoryAccounting")
178            .field("snapshot", &self.snapshot())
179            .finish()
180    }
181}
182
183impl MemoryAccounting {
184    pub fn new(budget: MemoryBudget) -> Self {
185        Self {
186            budget,
187            current: std::array::from_fn(|_| AtomicU64::new(0)),
188            peaks: std::array::from_fn(|_| AtomicU64::new(0)),
189            total_current: AtomicU64::new(0),
190            total_peak: AtomicU64::new(0),
191            soft_limit_events: AtomicU64::new(0),
192            hard_limit_rejections: AtomicU64::new(0),
193        }
194    }
195
196    pub fn budget(&self) -> MemoryBudget {
197        self.budget
198    }
199
200    /// Replace an observed usage value for a category.
201    pub fn report_usage(&self, category: MemoryCategory, bytes: u64) {
202        let slot = &self.current[category as usize];
203        let previous = slot.swap(bytes, Ordering::Relaxed);
204        let total = if bytes >= previous {
205            self.total_current
206                .fetch_add(bytes - previous, Ordering::Relaxed)
207                .saturating_add(bytes - previous)
208        } else {
209            self.total_current
210                .fetch_sub(previous - bytes, Ordering::Relaxed)
211                .saturating_sub(previous - bytes)
212        };
213        self.update_peak(&self.peaks[category as usize], bytes);
214        self.update_peak(&self.total_peak, total);
215        self.record_soft_crossing(total.saturating_sub(bytes), total);
216    }
217
218    /// Reserve temporary memory for a background operation.
219    pub fn try_reserve(
220        self: &Arc<Self>,
221        category: MemoryCategory,
222        bytes: u64,
223    ) -> StorageResult<MemoryReservation> {
224        let category_current = self.current[category as usize].load(Ordering::Relaxed);
225        if category_current.saturating_add(bytes) > self.budget.category_limit(category) {
226            self.hard_limit_rejections.fetch_add(1, Ordering::Relaxed);
227            return Err(StorageError::new(
228                StorageErrorKind::CapacityExceeded,
229                format!(
230                    "{} memory budget exceeded: requested {} bytes",
231                    category.name(),
232                    bytes
233                ),
234            ));
235        }
236
237        let mut current = self.total_current.load(Ordering::Relaxed);
238        loop {
239            let next = current.saturating_add(bytes);
240            if next > self.budget.hard_limit_bytes {
241                self.hard_limit_rejections.fetch_add(1, Ordering::Relaxed);
242                return Err(StorageError::new(
243                    StorageErrorKind::CapacityExceeded,
244                    format!(
245                        "total memory hard limit exceeded: requested {} bytes",
246                        bytes
247                    ),
248                ));
249            }
250            match self.total_current.compare_exchange_weak(
251                current,
252                next,
253                Ordering::AcqRel,
254                Ordering::Relaxed,
255            ) {
256                Ok(_) => {
257                    self.current[category as usize].fetch_add(bytes, Ordering::Relaxed);
258                    self.update_peak(&self.peaks[category as usize], category_current + bytes);
259                    self.update_peak(&self.total_peak, next);
260                    self.record_soft_crossing(current, next);
261                    return Ok(MemoryReservation {
262                        accounting: Arc::clone(self),
263                        category,
264                        bytes,
265                    });
266                }
267                Err(observed) => current = observed,
268            }
269        }
270    }
271
272    /// Reserve memory with automatic spilling on hard-limit failure.
273    ///
274    /// If `try_reserve` fails with `CapacityExceeded`, invokes `spill_fn` to
275    /// free memory, then retries the reservation. If the spill fails or the
276    /// retry also fails, the original error is returned.
277    pub fn try_reserve_with_spill<F>(
278        self: &Arc<Self>,
279        category: MemoryCategory,
280        bytes: u64,
281        spill_fn: F,
282    ) -> StorageResult<MemoryReservation>
283    where
284        F: Fn() -> Option<u64>,
285    {
286        match self.try_reserve(category, bytes) {
287            Ok(reservation) => Ok(reservation),
288            Err(original_err) => {
289                if let Some(freed) = spill_fn() {
290                    if freed > 0 {
291                        return self.try_reserve(category, bytes);
292                    }
293                }
294                Err(original_err)
295            }
296        }
297    }
298
299    /// Release memory from a category. Called by MemoryReservation::Drop
300    /// and by the spiller after evicting segments.
301    pub(crate) fn release(&self, category: MemoryCategory, bytes: u64) {
302        self.current[category as usize].fetch_sub(bytes, Ordering::Relaxed);
303        self.total_current.fetch_sub(bytes, Ordering::Relaxed);
304    }
305
306    /// Release memory from a category (public wrapper for eviction callbacks).
307    ///
308    /// This is the public interface for components like the cache eviction
309    /// listener to report memory that was freed outside the reservation system.
310    pub fn release_category(&self, category: MemoryCategory, bytes: u64) {
311        self.release(category, bytes);
312    }
313
314    pub fn snapshot(&self) -> ResourceSnapshot {
315        ResourceSnapshot {
316            budget: self.budget,
317            categories: std::array::from_fn(|index| MemoryUsage {
318                current_bytes: self.current[index].load(Ordering::Relaxed),
319                peak_bytes: self.peaks[index].load(Ordering::Relaxed),
320            }),
321            total_current_bytes: self.total_current.load(Ordering::Relaxed),
322            total_peak_bytes: self.total_peak.load(Ordering::Relaxed),
323            soft_limit_events: self.soft_limit_events.load(Ordering::Relaxed),
324            hard_limit_rejections: self.hard_limit_rejections.load(Ordering::Relaxed),
325            active_snapshots: 0,
326            oldest_snapshot_ts: Timestamp::MAX,
327            tombstone_count: 0,
328            tombstone_memory_bytes: 0,
329            vertex_shard_count: 0,
330        }
331    }
332
333    fn update_peak(&self, peak: &AtomicU64, value: u64) {
334        let mut previous = peak.load(Ordering::Relaxed);
335        while value > previous {
336            match peak.compare_exchange_weak(previous, value, Ordering::Relaxed, Ordering::Relaxed)
337            {
338                Ok(_) => break,
339                Err(observed) => previous = observed,
340            }
341        }
342    }
343
344    fn record_soft_crossing(&self, previous: u64, current: u64) {
345        if previous < self.budget.soft_limit_bytes && current >= self.budget.soft_limit_bytes {
346            self.soft_limit_events.fetch_add(1, Ordering::Relaxed);
347        }
348    }
349}
350
351/// RAII handle for temporary background memory.
352pub struct MemoryReservation {
353    accounting: Arc<MemoryAccounting>,
354    category: MemoryCategory,
355    bytes: u64,
356}
357
358impl std::fmt::Debug for MemoryReservation {
359    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
360        f.debug_struct("MemoryReservation")
361            .field("category", &self.category)
362            .field("bytes", &self.bytes)
363            .finish()
364    }
365}
366
367impl Drop for MemoryReservation {
368    fn drop(&mut self) {
369        self.accounting.release(self.category, self.bytes);
370    }
371}
372
373#[cfg(test)]
374mod tests {
375    use super::*;
376
377    #[test]
378    fn budget_rejects_invalid_limits() {
379        assert!(MemoryBudget::new(0, 1, 0.8, 0.95).is_err());
380        assert!(MemoryBudget::new(100, 101, 0.8, 0.95).is_err());
381        assert!(MemoryBudget::new(100, 20, 0.95, 0.8).is_err());
382    }
383
384    #[test]
385    fn accounting_tracks_pressure_and_releases_reservations() {
386        let budget = MemoryBudget::new(100, 20, 0.8, 0.95).expect("valid budget");
387        let accounting = Arc::new(MemoryAccounting::new(budget));
388
389        let reservation = accounting
390            .try_reserve(MemoryCategory::Background, 90)
391            .expect("reservation below hard limit");
392        let snapshot = accounting.snapshot();
393        assert_eq!(snapshot.total_current_bytes, 90);
394        assert!(snapshot.soft_limit_exceeded());
395        assert_eq!(snapshot.soft_limit_events, 1);
396
397        assert!(accounting
398            .try_reserve(MemoryCategory::Background, 6)
399            .is_err());
400        assert_eq!(accounting.snapshot().hard_limit_rejections, 1);
401
402        drop(reservation);
403        assert_eq!(accounting.snapshot().total_current_bytes, 0);
404    }
405
406    #[test]
407    fn index_budget_is_independent_from_total_budget() {
408        let budget = MemoryBudget::new(100, 20, 0.8, 0.95).expect("valid budget");
409        let accounting = Arc::new(MemoryAccounting::new(budget));
410        assert!(accounting.try_reserve(MemoryCategory::Index, 21).is_err());
411        assert!(accounting.try_reserve(MemoryCategory::Data, 20).is_ok());
412    }
413}
linkrs_storage/engine/resource_budget.rs

linkrs_storage/engine/
resource_budget.rs
