alopex_chirps/
memory.rs

1//! Node-wide logical memory accounting and runtime budget controls.
2//!
3//! The manager accounts for Chirps-owned allocations. It deliberately does not
4//! claim to be an RSS/cgroup sampler; those measurements belong to the v0.6.3
5//! stability evidence tooling.
6
7use crate::buffer::MessageBuffer;
8use std::collections::VecDeque;
9use std::sync::atomic::{AtomicUsize, Ordering};
10use std::sync::{Arc, Mutex};
11use thiserror::Error;
12
13const DEFAULT_TOTAL_BUDGET: usize = 256 * 1024 * 1024;
14const DEFAULT_MESSAGE_BUFFER_LIMIT: usize = 64 * 1024 * 1024;
15const DEFAULT_RAFT_LOG_CACHE_LIMIT: usize = 32 * 1024 * 1024;
16const DEFAULT_CONNECTION_POOL_LIMIT: usize = 16 * 1024 * 1024;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum MemoryComponent {
20    MessageBuffer,
21    RaftLogCache,
22    ConnectionPool,
23    BlockCache,
24}
25
26#[derive(Debug, Clone, Copy, PartialEq)]
27pub struct MemoryConfig {
28    pub total_budget: usize,
29    pub message_buffer_limit: usize,
30    pub raft_log_cache_limit: usize,
31    pub connection_pool_limit: usize,
32    pub backpressure_threshold: f32,
33    pub emergency_threshold: f32,
34}
35
36impl Default for MemoryConfig {
37    fn default() -> Self {
38        Self {
39            total_budget: DEFAULT_TOTAL_BUDGET,
40            message_buffer_limit: DEFAULT_MESSAGE_BUFFER_LIMIT,
41            raft_log_cache_limit: DEFAULT_RAFT_LOG_CACHE_LIMIT,
42            connection_pool_limit: DEFAULT_CONNECTION_POOL_LIMIT,
43            backpressure_threshold: 0.80,
44            emergency_threshold: 0.95,
45        }
46    }
47}
48
49#[derive(Debug, Error, PartialEq)]
50pub enum MemoryError {
51    #[error("memory total_budget must be greater than zero")]
52    ZeroBudget,
53    #[error("memory threshold must be finite and between zero and one")]
54    InvalidThreshold,
55    #[error("memory component limit exceeds total budget")]
56    ComponentLimitExceedsBudget,
57    #[error("at least one memory sample is required")]
58    EmptySamples,
59    #[error("allocation ratios must be finite, non-negative, and sum to one")]
60    InvalidAllocationRatio,
61}
62
63#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
64pub struct MemoryStats {
65    pub total_budget: usize,
66    pub current_usage: usize,
67    pub message_buffer_bytes: usize,
68    pub raft_log_cache_bytes: usize,
69    pub connection_pool_bytes: usize,
70    pub block_cache_bytes: usize,
71    pub budget_exceeded: bool,
72}
73
74/// Dynamic shares of the node-wide memory budget.
75#[derive(Debug, Clone, Copy, PartialEq)]
76pub struct AllocationRatio {
77    pub message_buffer: f32,
78    pub raft_cache: f32,
79    pub connection_pool: f32,
80    pub block_cache: f32,
81}
82
83impl Default for AllocationRatio {
84    fn default() -> Self {
85        Self {
86            message_buffer: 0.30,
87            raft_cache: 0.20,
88            connection_pool: 0.10,
89            block_cache: 0.40,
90        }
91    }
92}
93
94impl AllocationRatio {
95    fn validate(self) -> Result<(), MemoryError> {
96        let values = [
97            self.message_buffer,
98            self.raft_cache,
99            self.connection_pool,
100            self.block_cache,
101        ];
102        if values
103            .iter()
104            .any(|value| !value.is_finite() || *value < 0.0)
105            || (values.iter().sum::<f32>() - 1.0).abs() > 0.0001
106        {
107            return Err(MemoryError::InvalidAllocationRatio);
108        }
109        Ok(())
110    }
111}
112
113/// Workload weights used by [`IntegratedCacheManager::rebalance`].
114pub type WorkloadProfile = AllocationRatio;
115
116/// A bounded LRU cache for recently used Raft log payloads.
117#[derive(Debug)]
118pub struct RaftLogCache {
119    max_bytes: usize,
120    used_bytes: usize,
121    entries: VecDeque<(u64, Vec<u8>)>,
122}
123
124impl RaftLogCache {
125    pub fn new(max_bytes: usize) -> Self {
126        Self {
127            max_bytes,
128            used_bytes: 0,
129            entries: VecDeque::new(),
130        }
131    }
132
133    pub fn insert(&mut self, index: u64, payload: Vec<u8>) {
134        if let Some(position) = self.entries.iter().position(|(key, _)| *key == index)
135            && let Some((_, old)) = self.entries.remove(position)
136        {
137            self.used_bytes = self.used_bytes.saturating_sub(old.len());
138        }
139        if payload.len() > self.max_bytes {
140            return;
141        }
142        while self.used_bytes.saturating_add(payload.len()) > self.max_bytes {
143            let Some((_, evicted)) = self.entries.pop_front() else {
144                break;
145            };
146            self.used_bytes = self.used_bytes.saturating_sub(evicted.len());
147        }
148        self.used_bytes = self.used_bytes.saturating_add(payload.len());
149        self.entries.push_back((index, payload));
150    }
151
152    pub fn get(&mut self, index: u64) -> Option<Vec<u8>> {
153        let position = self.entries.iter().position(|(key, _)| *key == index)?;
154        let entry = self.entries.remove(position)?;
155        let payload = entry.1.clone();
156        self.entries.push_back(entry);
157        Some(payload)
158    }
159
160    pub fn used_bytes(&self) -> usize {
161        self.used_bytes
162    }
163
164    fn evict_bytes(&mut self, target_bytes: usize) -> usize {
165        let mut evicted = 0;
166        while evicted < target_bytes {
167            let Some((_, payload)) = self.entries.pop_front() else {
168                break;
169            };
170            evicted = evicted.saturating_add(payload.len());
171            self.used_bytes = self.used_bytes.saturating_sub(payload.len());
172        }
173        evicted
174    }
175}
176
177/// Usage adapter for alopex-core's block cache.
178///
179/// alopex-core 0.3 does not expose a public BlockCache type. This adapter is
180/// intentionally limited to the accounting/eviction contract and can wrap a
181/// concrete core cache when that API becomes available without changing the
182/// Chirps wire or persistence formats.
183#[derive(Debug)]
184pub struct BlockCacheHandle {
185    capacity: AtomicUsize,
186    used_bytes: AtomicUsize,
187}
188
189impl BlockCacheHandle {
190    pub fn new(capacity: usize) -> Self {
191        Self {
192            capacity: AtomicUsize::new(capacity),
193            used_bytes: AtomicUsize::new(0),
194        }
195    }
196
197    pub fn set_capacity(&self, capacity: usize) {
198        self.capacity.store(capacity, Ordering::Relaxed);
199    }
200
201    pub fn capacity(&self) -> usize {
202        self.capacity.load(Ordering::Relaxed)
203    }
204
205    pub fn set_used_bytes(&self, bytes: usize) {
206        self.used_bytes.store(bytes, Ordering::Relaxed);
207    }
208
209    pub fn used_bytes(&self) -> usize {
210        self.used_bytes.load(Ordering::Relaxed)
211    }
212
213    fn evict_bytes(&self, target_bytes: usize) -> usize {
214        loop {
215            let current = self.used_bytes();
216            let evicted = current.min(target_bytes);
217            if self
218                .used_bytes
219                .compare_exchange(
220                    current,
221                    current - evicted,
222                    Ordering::Relaxed,
223                    Ordering::Relaxed,
224                )
225                .is_ok()
226            {
227                return evicted;
228            }
229        }
230    }
231}
232
233#[derive(Debug, Clone, Copy, PartialEq, Eq)]
234pub struct UnifiedMemoryMetrics {
235    pub total_budget: usize,
236    pub current_usage: usize,
237    pub message_buffer_bytes: usize,
238    pub raft_cache_bytes: usize,
239    pub connection_pool_bytes: usize,
240    pub block_cache_bytes: usize,
241    pub evicted_bytes: usize,
242}
243
244/// One best-effort process/cgroup memory observation.
245#[derive(Debug, Clone, Copy, PartialEq, Eq)]
246pub struct MemoryMeasurement {
247    pub rss_bytes: Option<u64>,
248    pub cgroup_current_bytes: Option<u64>,
249    pub cgroup_limit_bytes: Option<u64>,
250}
251
252impl MemoryMeasurement {
253    pub const fn from_values(
254        rss_bytes: Option<u64>,
255        cgroup_current_bytes: Option<u64>,
256        cgroup_limit_bytes: Option<u64>,
257    ) -> Self {
258        Self {
259            rss_bytes,
260            cgroup_current_bytes,
261            cgroup_limit_bytes,
262        }
263    }
264
265    /// Selects RSS when available and otherwise falls back to cgroup usage.
266    pub fn observed_bytes(self) -> Option<u64> {
267        self.rss_bytes.or(self.cgroup_current_bytes)
268    }
269
270    /// Captures the process RSS and common cgroup v2/v1 files when present.
271    pub fn capture() -> Self {
272        let rss_bytes = std::fs::read_to_string("/proc/self/status")
273            .ok()
274            .and_then(|status| {
275                status.lines().find_map(|line| {
276                    let value = line.strip_prefix("VmRSS:")?.split_whitespace().next()?;
277                    value
278                        .parse::<u64>()
279                        .ok()
280                        .map(|kib| kib.saturating_mul(1024))
281                })
282            });
283        let cgroup_current_bytes = read_first_number(&[
284            "/sys/fs/cgroup/memory.current",
285            "/sys/fs/cgroup/memory/memory.usage_in_bytes",
286        ]);
287        let cgroup_limit_bytes = read_first_number(&[
288            "/sys/fs/cgroup/memory.max",
289            "/sys/fs/cgroup/memory/memory.limit_in_bytes",
290        ]);
291        Self::from_values(rss_bytes, cgroup_current_bytes, cgroup_limit_bytes)
292    }
293}
294
295/// Summary used by the high-load memory stability measurement contract.
296#[derive(Debug, Clone, Copy, PartialEq, Eq)]
297pub struct MemoryStabilityReport {
298    pub peak_bytes: u64,
299    pub final_bytes: u64,
300    pub growth_bytes: u64,
301    pub within_budget: bool,
302    pub stable: bool,
303}
304
305impl MemoryStabilityReport {
306    pub fn from_samples(
307        samples: &[MemoryMeasurement],
308        budget_bytes: u64,
309        allowed_growth_bytes: u64,
310    ) -> Result<Self, MemoryError> {
311        let values: Vec<u64> = samples
312            .iter()
313            .copied()
314            .filter_map(MemoryMeasurement::observed_bytes)
315            .collect();
316        let Some((&first, rest)) = values.split_first() else {
317            return Err(MemoryError::EmptySamples);
318        };
319        let peak_bytes = values.iter().copied().max().unwrap_or(first);
320        let final_bytes = values.last().copied().unwrap_or(first);
321        let growth_bytes = final_bytes.saturating_sub(first);
322        Ok(Self {
323            peak_bytes,
324            final_bytes,
325            growth_bytes,
326            within_budget: peak_bytes <= budget_bytes,
327            stable: growth_bytes <= allowed_growth_bytes
328                && rest
329                    .iter()
330                    .all(|value| value.saturating_sub(first) <= allowed_growth_bytes),
331        })
332    }
333}
334
335fn read_first_number(paths: &[&str]) -> Option<u64> {
336    paths.iter().find_map(|path| {
337        let value = std::fs::read_to_string(path).ok()?.trim().to_owned();
338        if value == "max" {
339            None
340        } else {
341            value.parse().ok()
342        }
343    })
344}
345
346#[derive(Debug)]
347pub struct MemoryManager {
348    state: Mutex<MemoryState>,
349}
350
351#[derive(Debug)]
352struct MemoryState {
353    config: MemoryConfig,
354    stats: MemoryStats,
355}
356
357impl MemoryManager {
358    pub fn new(config: MemoryConfig) -> Result<Self, MemoryError> {
359        validate_config(&config)?;
360        Ok(Self {
361            state: Mutex::new(MemoryState {
362                stats: MemoryStats {
363                    total_budget: config.total_budget,
364                    ..MemoryStats::default()
365                },
366                config,
367            }),
368        })
369    }
370
371    pub fn resize_memory_budget(&self, new_budget: usize) -> Result<(), MemoryError> {
372        if new_budget == 0 {
373            return Err(MemoryError::ZeroBudget);
374        }
375        let mut state = self.state.lock().expect("memory manager lock poisoned");
376        state.config.total_budget = new_budget;
377        state.stats.total_budget = new_budget;
378        refresh_budget_flag(&mut state.stats);
379        Ok(())
380    }
381
382    pub fn set_component_usage(&self, component: MemoryComponent, bytes: usize) {
383        let mut state = self.state.lock().expect("memory manager lock poisoned");
384        let slot = match component {
385            MemoryComponent::MessageBuffer => &mut state.stats.message_buffer_bytes,
386            MemoryComponent::RaftLogCache => &mut state.stats.raft_log_cache_bytes,
387            MemoryComponent::ConnectionPool => &mut state.stats.connection_pool_bytes,
388            MemoryComponent::BlockCache => &mut state.stats.block_cache_bytes,
389        };
390        *slot = bytes;
391        refresh_budget_flag(&mut state.stats);
392    }
393
394    fn set_component_limit(&self, component: MemoryComponent, bytes: usize) {
395        let mut state = self.state.lock().expect("memory manager lock poisoned");
396        let limit = bytes.min(state.config.total_budget);
397        match component {
398            MemoryComponent::MessageBuffer => state.config.message_buffer_limit = limit,
399            MemoryComponent::RaftLogCache => state.config.raft_log_cache_limit = limit,
400            MemoryComponent::ConnectionPool => state.config.connection_pool_limit = limit,
401            MemoryComponent::BlockCache => {}
402        }
403    }
404
405    /// Applies configured component caps and gives any remaining budget to the
406    /// block cache. This is the explicit, synchronous eviction boundary exposed
407    /// to operators; transport/WAL owners can call it before releasing storage.
408    pub fn trigger_gc(&self) -> Result<(), MemoryError> {
409        let mut state = self.state.lock().expect("memory manager lock poisoned");
410        state.stats.message_buffer_bytes = state
411            .stats
412            .message_buffer_bytes
413            .min(state.config.message_buffer_limit);
414        state.stats.raft_log_cache_bytes = state
415            .stats
416            .raft_log_cache_bytes
417            .min(state.config.raft_log_cache_limit);
418        state.stats.connection_pool_bytes = state
419            .stats
420            .connection_pool_bytes
421            .min(state.config.connection_pool_limit);
422        let fixed = state
423            .stats
424            .message_buffer_bytes
425            .saturating_add(state.stats.raft_log_cache_bytes)
426            .saturating_add(state.stats.connection_pool_bytes);
427        let remaining = state.config.total_budget.saturating_sub(fixed);
428        state.stats.block_cache_bytes = state.stats.block_cache_bytes.min(remaining);
429        refresh_budget_flag(&mut state.stats);
430        Ok(())
431    }
432
433    pub fn get_memory_stats(&self) -> MemoryStats {
434        self.state
435            .lock()
436            .expect("memory manager lock poisoned")
437            .stats
438    }
439}
440
441/// Coordinates the memory owned by Chirps subsystems under the existing
442/// [`MemoryManager`] budget.
443pub struct IntegratedCacheManager {
444    pub message_buffer: MessageBuffer,
445    pub raft_cache: RaftLogCache,
446    pub block_cache: Arc<BlockCacheHandle>,
447    pub total_budget: usize,
448    pub allocation_ratio: AllocationRatio,
449    memory: Arc<MemoryManager>,
450    connection_pool_bytes: usize,
451    evicted_bytes: usize,
452}
453
454impl IntegratedCacheManager {
455    pub fn new(config: MemoryConfig) -> Result<Self, MemoryError> {
456        Self::with_block_cache(
457            config,
458            Arc::new(BlockCacheHandle::new(
459                (config.total_budget as f32 * AllocationRatio::default().block_cache) as usize,
460            )),
461        )
462    }
463
464    pub fn with_block_cache(
465        config: MemoryConfig,
466        block_cache: Arc<BlockCacheHandle>,
467    ) -> Result<Self, MemoryError> {
468        let memory = Arc::new(MemoryManager::new(config)?);
469        let allocation_ratio = AllocationRatio::default();
470        allocation_ratio.validate()?;
471        let manager = Self {
472            message_buffer: MessageBuffer::new(
473                config.message_buffer_limit,
474                config.backpressure_threshold,
475                config.emergency_threshold,
476            ),
477            raft_cache: RaftLogCache::new(config.raft_log_cache_limit),
478            block_cache,
479            total_budget: config.total_budget,
480            allocation_ratio,
481            memory,
482            connection_pool_bytes: 0,
483            evicted_bytes: 0,
484        };
485        manager.sync_usage();
486        Ok(manager)
487    }
488
489    /// Reallocates component caps and evicts immediately when a new cap is
490    /// below current usage. The shared MemoryManager remains the accounting
491    /// source of truth used by MeshHandle.
492    pub fn rebalance(&mut self, workload: WorkloadProfile) {
493        if workload.validate().is_err() {
494            return;
495        }
496        self.allocation_ratio = workload;
497        let budget = self.total_budget;
498        let message_limit = (budget as f32 * workload.message_buffer) as usize;
499        let raft_limit = (budget as f32 * workload.raft_cache) as usize;
500        let connection_limit = (budget as f32 * workload.connection_pool) as usize;
501        let block_limit = (budget as f32 * workload.block_cache) as usize;
502        self.memory
503            .set_component_limit(MemoryComponent::MessageBuffer, message_limit);
504        self.memory
505            .set_component_limit(MemoryComponent::RaftLogCache, raft_limit);
506        self.memory
507            .set_component_limit(MemoryComponent::ConnectionPool, connection_limit);
508        self.block_cache.set_capacity(block_limit);
509        self.evicted_bytes = self.evicted_bytes.saturating_add(
510            self.evict_message_bytes(
511                self.message_buffer
512                    .used_bytes()
513                    .saturating_sub(message_limit),
514            ),
515        );
516        self.evicted_bytes = self.evicted_bytes.saturating_add(
517            self.raft_cache
518                .evict_bytes(self.raft_cache.used_bytes().saturating_sub(raft_limit)),
519        );
520        self.evicted_bytes = self.evicted_bytes.saturating_add(
521            self.block_cache
522                .evict_bytes(self.block_cache.used_bytes().saturating_sub(block_limit)),
523        );
524        self.sync_usage();
525    }
526
527    /// Frees at least as much as possible from the least durable in-memory
528    /// layers and returns the number of bytes actually released.
529    pub fn emergency_evict(&mut self, target_bytes: usize) -> usize {
530        let mut remaining = target_bytes;
531        let mut evicted = 0;
532        let released = self.block_cache.evict_bytes(remaining);
533        remaining = remaining.saturating_sub(released);
534        evicted += released;
535        if remaining > 0 {
536            let released = self.raft_cache.evict_bytes(remaining);
537            remaining = remaining.saturating_sub(released);
538            evicted += released;
539        }
540        if remaining > 0 {
541            evicted += self.evict_message_bytes(remaining);
542        }
543        self.evicted_bytes = self.evicted_bytes.saturating_add(evicted);
544        self.sync_usage();
545        evicted
546    }
547
548    pub fn get_unified_metrics(&self) -> UnifiedMemoryMetrics {
549        self.sync_usage();
550        let stats = self.memory.get_memory_stats();
551        UnifiedMemoryMetrics {
552            total_budget: stats.total_budget,
553            current_usage: stats.current_usage,
554            message_buffer_bytes: stats.message_buffer_bytes,
555            raft_cache_bytes: stats.raft_log_cache_bytes,
556            connection_pool_bytes: stats.connection_pool_bytes,
557            block_cache_bytes: stats.block_cache_bytes,
558            evicted_bytes: self.evicted_bytes,
559        }
560    }
561
562    /// Records bytes owned by the QUIC connection/session layer in the same
563    /// budget. The transport can update this at its own sampling boundary.
564    pub fn set_connection_pool_usage(&mut self, bytes: usize) {
565        self.connection_pool_bytes = bytes;
566        self.sync_usage();
567    }
568
569    pub fn memory_manager(&self) -> Arc<MemoryManager> {
570        Arc::clone(&self.memory)
571    }
572
573    fn evict_message_bytes(&mut self, target_bytes: usize) -> usize {
574        let mut evicted = 0;
575        while evicted < target_bytes {
576            let Some(message) = self.message_buffer.pop() else {
577                break;
578            };
579            evicted = evicted.saturating_add(message.payload.len());
580        }
581        evicted
582    }
583
584    fn sync_usage(&self) {
585        self.memory.set_component_usage(
586            MemoryComponent::MessageBuffer,
587            self.message_buffer.used_bytes(),
588        );
589        self.memory
590            .set_component_usage(MemoryComponent::RaftLogCache, self.raft_cache.used_bytes());
591        self.memory
592            .set_component_usage(MemoryComponent::ConnectionPool, self.connection_pool_bytes);
593        self.memory
594            .set_component_usage(MemoryComponent::BlockCache, self.block_cache.used_bytes());
595    }
596}
597
598fn validate_config(config: &MemoryConfig) -> Result<(), MemoryError> {
599    if config.total_budget == 0 {
600        return Err(MemoryError::ZeroBudget);
601    }
602    if config.message_buffer_limit > config.total_budget
603        || config.raft_log_cache_limit > config.total_budget
604        || config.connection_pool_limit > config.total_budget
605    {
606        return Err(MemoryError::ComponentLimitExceedsBudget);
607    }
608    if !config.backpressure_threshold.is_finite()
609        || !config.emergency_threshold.is_finite()
610        || !(0.0..=1.0).contains(&config.backpressure_threshold)
611        || !(0.0..=1.0).contains(&config.emergency_threshold)
612        || config.backpressure_threshold > config.emergency_threshold
613    {
614        return Err(MemoryError::InvalidThreshold);
615    }
616    Ok(())
617}
618
619fn refresh_budget_flag(stats: &mut MemoryStats) {
620    stats.current_usage = stats
621        .message_buffer_bytes
622        .saturating_add(stats.raft_log_cache_bytes)
623        .saturating_add(stats.connection_pool_bytes)
624        .saturating_add(stats.block_cache_bytes);
625    stats.budget_exceeded = stats.current_usage > stats.total_budget;
626}
alopex_chirps/memory.rs

alopex_chirps/
memory.rs
