astra_core/
memory.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use crate::errors::StoreError;
4
5#[derive(Debug)]
6pub struct MemoryTracker {
7    max_bytes: usize,
8    used_bytes: AtomicUsize,
9}
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum MemoryPressure {
13    Normal,
14    High,
15    Critical,
16}
17
18impl MemoryTracker {
19    pub fn new(max_bytes: usize) -> Self {
20        Self {
21            max_bytes,
22            used_bytes: AtomicUsize::new(0),
23        }
24    }
25
26    pub fn used_bytes(&self) -> usize {
27        self.used_bytes.load(Ordering::Relaxed)
28    }
29
30    pub fn max_bytes(&self) -> usize {
31        self.max_bytes
32    }
33
34    pub fn pressure(&self) -> MemoryPressure {
35        let used = self.used_bytes();
36        let ratio = used as f64 / self.max_bytes.max(1) as f64;
37        if ratio >= 0.95 {
38            MemoryPressure::Critical
39        } else if ratio >= 0.85 {
40            MemoryPressure::High
41        } else {
42            MemoryPressure::Normal
43        }
44    }
45
46    pub fn try_increase(&self, amount: usize) -> Result<(), StoreError> {
47        if amount == 0 {
48            return Ok(());
49        }
50
51        let mut prev = self.used_bytes.load(Ordering::Relaxed);
52        loop {
53            let next = prev.saturating_add(amount);
54            if next > self.max_bytes {
55                return Err(StoreError::ResourceExhausted(format!(
56                    "memory cap exceeded: requested={} used={} max={}",
57                    amount, prev, self.max_bytes
58                )));
59            }
60
61            match self.used_bytes.compare_exchange_weak(
62                prev,
63                next,
64                Ordering::SeqCst,
65                Ordering::Relaxed,
66            ) {
67                Ok(_) => break,
68                Err(actual) => prev = actual,
69            }
70        }
71        Ok(())
72    }
73
74    pub fn decrease(&self, amount: usize) {
75        if amount == 0 {
76            return;
77        }
78
79        let mut prev = self.used_bytes.load(Ordering::Relaxed);
80        loop {
81            let next = prev.saturating_sub(amount);
82            match self.used_bytes.compare_exchange_weak(
83                prev,
84                next,
85                Ordering::SeqCst,
86                Ordering::Relaxed,
87            ) {
88                Ok(_) => break,
89                Err(actual) => prev = actual,
90            }
91        }
92    }
93}
astra_core/memory.rs

astra_core/
memory.rs
