tale_ndjson/
memory_budget.rs

1//! Memory budget management for adaptive chunking
2//!
3//! This module provides comprehensive memory budget tracking and enforcement
4//! for the chunked file processing system. It monitors system memory usage,
5//! enforces per-process limits, and provides pressure-based adaptation signals.
6//!
7//! ## Key Features
8//! - **Global memory budget**: System-wide memory limit enforcement
9//! - **Per-reader tracking**: Individual memory usage monitoring
10//! - **Pressure detection**: Multi-level memory pressure signals
11//! - **Adaptive response**: Automatic chunk size and buffer adjustments
12//! - **Graceful degradation**: Fallback strategies under memory pressure
13//!
14//! ## Usage
15//! ```no_run
16//! use tale_ndjson::{MemoryBudget, MemoryPressure};
17//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
18//! let mut budget = MemoryBudget::new(100 * 1024 * 1024)?; // 100MB limit
19//!
20//! // Allocate memory for a chunk
21//! let chunk_size = 4096;
22//! if let Some(allocation) = budget.try_allocate(chunk_size, "reader_1")? {
23//!     // Process with allocated memory
24//!     allocation.deallocate(); // Automatic on drop
25//! }
26//!
27//! // Check memory pressure
28//! match budget.current_pressure()? {
29//!     MemoryPressure::Low => { /* normal operation */ },
30//!     MemoryPressure::Moderate => { /* mild optimization */ },
31//!     MemoryPressure::High => { /* reduce chunk sizes */ },
32//!     MemoryPressure::Critical => { /* emergency measures */ },
33//! }
34//! # Ok(())
35//! # }
36//! ```
37
38use std::collections::HashMap;
39use std::sync::{Arc, RwLock};
40use std::time::{Duration, Instant};
41
42use crate::errors::TaleError;
43
44/// Memory pressure levels for adaptive response
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum MemoryPressure {
47    /// Memory usage < 60% of limit - normal operation
48    Low,
49    /// Memory usage 60-85% of limit - start reducing allocations
50    Moderate,
51    /// Memory usage 85-95% of limit - aggressive reduction
52    High,
53    /// Memory usage > 95% of limit - emergency measures
54    Critical,
55}
56
57impl MemoryPressure {
58    /// Get the adaptation factor for chunk sizes based on pressure level
59    pub fn chunk_size_factor(&self) -> f64 {
60        match self {
61            MemoryPressure::Low => 1.0,       // No reduction
62            MemoryPressure::Moderate => 0.8,  // 20% reduction
63            MemoryPressure::High => 0.5,      // 50% reduction
64            MemoryPressure::Critical => 0.25, // 75% reduction
65        }
66    }
67
68    /// Whether to enable low-memory optimizations
69    pub fn requires_optimization(&self) -> bool {
70        matches!(self, MemoryPressure::High | MemoryPressure::Critical)
71    }
72}
73
74/// Individual memory allocation tracking
75#[derive(Debug)]
76pub struct MemoryAllocation {
77    size: usize,
78    reader_id: String,
79    allocated_at: Instant,
80    budget: Arc<RwLock<MemoryBudgetInner>>,
81}
82
83impl MemoryAllocation {
84    fn new(size: usize, reader_id: String, budget: Arc<RwLock<MemoryBudgetInner>>) -> Self {
85        Self {
86            size,
87            reader_id,
88            allocated_at: Instant::now(),
89            budget,
90        }
91    }
92
93    /// Get the size of this allocation
94    pub fn size(&self) -> usize {
95        self.size
96    }
97
98    /// Get how long this allocation has been active
99    pub fn age(&self) -> Duration {
100        self.allocated_at.elapsed()
101    }
102
103    /// Manually deallocate this memory (automatic on drop)
104    pub fn deallocate(self) {
105        drop(self); // Explicit drop to trigger deallocation
106    }
107}
108
109impl Drop for MemoryAllocation {
110    fn drop(&mut self) {
111        if let Ok(mut budget) = self.budget.write() {
112            budget.deallocate(self.size, &self.reader_id);
113        }
114    }
115}
116
117/// Per-reader memory usage statistics
118#[derive(Debug, Clone, Default)]
119pub struct ReaderMemoryStats {
120    /// Total bytes currently allocated
121    pub current_allocation: usize,
122    /// Peak allocation seen
123    pub peak_allocation: usize,
124    /// Number of active allocations
125    pub allocation_count: usize,
126    /// Number of times allocation failed
127    pub allocation_failures: usize,
128    /// Total allocations over lifetime
129    pub total_allocations: usize,
130}
131
132/// Internal budget state
133#[derive(Debug)]
134struct MemoryBudgetInner {
135    /// Maximum total memory allowed (in bytes)
136    total_limit: usize,
137    /// Currently allocated memory (in bytes)
138    current_usage: usize,
139    /// Peak memory usage seen
140    peak_usage: usize,
141    /// Per-reader memory tracking
142    reader_stats: HashMap<String, ReaderMemoryStats>,
143    /// System memory monitoring
144    last_system_check: Instant,
145    system_memory_available: usize,
146}
147
148impl MemoryBudgetInner {
149    fn new(total_limit: usize) -> Result<Self, TaleError> {
150        let system_memory = get_system_memory_available()?;
151
152        Ok(Self {
153            total_limit,
154            current_usage: 0,
155            peak_usage: 0,
156            reader_stats: HashMap::new(),
157            last_system_check: Instant::now(),
158            system_memory_available: system_memory,
159        })
160    }
161
162    fn try_allocate(&mut self, size: usize, reader_id: &str) -> Result<bool, TaleError> {
163        // Update system memory if it's been a while
164        if self.last_system_check.elapsed() > Duration::from_secs(1) {
165            self.system_memory_available = get_system_memory_available()?;
166            self.last_system_check = Instant::now();
167        }
168
169        // Check if allocation would exceed budget
170        let new_usage = self.current_usage + size;
171        if new_usage > self.total_limit {
172            // Update failure stats
173            let stats = self.reader_stats.entry(reader_id.to_string()).or_default();
174            stats.allocation_failures += 1;
175            return Ok(false);
176        }
177
178        // Check system memory availability (safety margin)
179        let system_safety_margin = self.system_memory_available / 4; // Keep 25% free
180        if size > system_safety_margin {
181            let stats = self.reader_stats.entry(reader_id.to_string()).or_default();
182            stats.allocation_failures += 1;
183            return Ok(false);
184        }
185
186        // Allocation successful
187        self.current_usage = new_usage;
188        self.peak_usage = self.peak_usage.max(new_usage);
189
190        // Update reader stats
191        let stats = self.reader_stats.entry(reader_id.to_string()).or_default();
192        stats.current_allocation += size;
193        stats.peak_allocation = stats.peak_allocation.max(stats.current_allocation);
194        stats.allocation_count += 1;
195        stats.total_allocations += 1;
196
197        Ok(true)
198    }
199
200    fn deallocate(&mut self, size: usize, reader_id: &str) {
201        self.current_usage = self.current_usage.saturating_sub(size);
202
203        if let Some(stats) = self.reader_stats.get_mut(reader_id) {
204            stats.current_allocation = stats.current_allocation.saturating_sub(size);
205            stats.allocation_count = stats.allocation_count.saturating_sub(1);
206        }
207    }
208
209    fn current_pressure(&self) -> MemoryPressure {
210        use crate::defaults::SystemDefaults;
211        let usage_ratio = self.current_usage as f64 / self.total_limit as f64;
212
213        match usage_ratio {
214            r if r < SystemDefaults::MEMORY_PRESSURE_LOW_THRESHOLD => MemoryPressure::Low,
215            r if r < SystemDefaults::MEMORY_PRESSURE_MODERATE_THRESHOLD => MemoryPressure::Moderate,
216            r if r < SystemDefaults::MEMORY_PRESSURE_HIGH_THRESHOLD => MemoryPressure::High,
217            _ => MemoryPressure::Critical,
218        }
219    }
220}
221
222/// Global memory budget manager for chunked file processing
223#[derive(Debug, Clone)]
224pub struct MemoryBudget {
225    inner: Arc<RwLock<MemoryBudgetInner>>,
226}
227
228impl MemoryBudget {
229    /// Create a new memory budget with the specified limit
230    pub fn new(total_limit: usize) -> Result<Self, TaleError> {
231        let inner = Arc::new(RwLock::new(MemoryBudgetInner::new(total_limit)?));
232        Ok(Self { inner })
233    }
234
235    /// Create a memory budget based on system memory
236    pub fn from_system_memory(percentage: f64) -> Result<Self, TaleError> {
237        let system_memory = get_system_memory_available()?;
238        let limit = (system_memory as f64 * percentage / 100.0) as usize;
239        Self::new(limit)
240    }
241
242    /// Try to allocate memory for a specific reader
243    pub fn try_allocate(&self, size: usize, reader_id: &str) -> Result<Option<MemoryAllocation>, TaleError> {
244        let mut inner = self
245            .inner
246            .write()
247            .map_err(|_| TaleError::MemoryError("Failed to acquire budget lock for allocation".to_string()))?;
248
249        if inner.try_allocate(size, reader_id)? {
250            let allocation = MemoryAllocation::new(size, reader_id.to_string(), self.inner.clone());
251            Ok(Some(allocation))
252        } else {
253            Ok(None)
254        }
255    }
256
257    /// Get current memory pressure level
258    pub fn current_pressure(&self) -> Result<MemoryPressure, TaleError> {
259        let inner = self
260            .inner
261            .read()
262            .map_err(|_| TaleError::MemoryError("Failed to acquire budget lock for pressure check".to_string()))?;
263        Ok(inner.current_pressure())
264    }
265
266    /// Get current memory usage statistics
267    pub fn usage_stats(&self) -> Result<MemoryBudgetStats, TaleError> {
268        let inner = self
269            .inner
270            .read()
271            .map_err(|_| TaleError::MemoryError("Failed to acquire budget lock for stats".to_string()))?;
272
273        Ok(MemoryBudgetStats {
274            total_limit: inner.total_limit,
275            current_usage: inner.current_usage,
276            peak_usage: inner.peak_usage,
277            pressure: inner.current_pressure(),
278            reader_count: inner.reader_stats.len(),
279            system_memory_available: inner.system_memory_available,
280        })
281    }
282
283    /// Get memory statistics for a specific reader
284    pub fn reader_stats(&self, reader_id: &str) -> Result<Option<ReaderMemoryStats>, TaleError> {
285        let inner = self
286            .inner
287            .read()
288            .map_err(|_| TaleError::MemoryError("Failed to acquire budget lock for reader stats".to_string()))?;
289        Ok(inner.reader_stats.get(reader_id).cloned())
290    }
291
292    /// Get recommended chunk size based on current memory pressure
293    pub fn recommended_chunk_size(&self, base_size: usize) -> Result<usize, TaleError> {
294        let pressure = self.current_pressure()?;
295        let factor = pressure.chunk_size_factor();
296        Ok((base_size as f64 * factor) as usize)
297    }
298
299    /// Check if emergency measures should be taken
300    pub fn requires_emergency_measures(&self) -> Result<bool, TaleError> {
301        let pressure = self.current_pressure()?;
302        Ok(matches!(pressure, MemoryPressure::Critical))
303    }
304}
305
306/// Memory budget usage statistics
307#[derive(Debug, Clone)]
308pub struct MemoryBudgetStats {
309    /// Total memory limit
310    pub total_limit: usize,
311    /// Currently used memory
312    pub current_usage: usize,
313    /// Peak memory usage
314    pub peak_usage: usize,
315    /// Current memory pressure level
316    pub pressure: MemoryPressure,
317    /// Number of active readers
318    pub reader_count: usize,
319    /// Available system memory
320    pub system_memory_available: usize,
321}
322
323impl MemoryBudgetStats {
324    /// Get memory usage as a percentage
325    pub fn usage_percentage(&self) -> f64 {
326        if self.total_limit > 0 {
327            (self.current_usage as f64 / self.total_limit as f64) * 100.0
328        } else {
329            0.0
330        }
331    }
332
333    /// Get available memory
334    pub fn available_memory(&self) -> usize {
335        self.total_limit.saturating_sub(self.current_usage)
336    }
337
338    /// Print a formatted report
339    pub fn print_report(&self) {
340        println!("Memory Budget Report:");
341        println!("====================");
342        println!("Total Limit:    {} MB", self.total_limit / (1024 * 1024));
343        println!(
344            "Current Usage:  {} MB ({:.1}%)",
345            self.current_usage / (1024 * 1024),
346            self.usage_percentage()
347        );
348        println!("Peak Usage:     {} MB", self.peak_usage / (1024 * 1024));
349        println!("Available:      {} MB", self.available_memory() / (1024 * 1024));
350        println!("Pressure Level: {:?}", self.pressure);
351        println!("Active Readers: {}", self.reader_count);
352        println!("System Memory:  {} MB", self.system_memory_available / (1024 * 1024));
353    }
354}
355
356/// Get available system memory
357fn get_system_memory_available() -> Result<usize, TaleError> {
358    // Try to get actual system memory stats
359    if let Some(stats) = memory_stats::memory_stats() {
360        // Use physical memory as a proxy for available memory
361        // This is conservative but safe
362        Ok(stats.physical_mem)
363    } else {
364        // Fallback: assume 1GB available (very conservative)
365        Ok(1024 * 1024 * 1024)
366    }
367}
368
369#[cfg(test)]
370mod tests {
371    use super::*;
372
373    #[test]
374    fn _memory_pressure_levels_work() {
375        assert_eq!(MemoryPressure::Low.chunk_size_factor(), 1.0);
376        assert_eq!(MemoryPressure::Moderate.chunk_size_factor(), 0.8);
377        assert_eq!(MemoryPressure::High.chunk_size_factor(), 0.5);
378        assert_eq!(MemoryPressure::Critical.chunk_size_factor(), 0.25);
379    }
380
381    #[test]
382    fn memory_budget_allocation_works() -> Result<(), TaleError> {
383        let budget = MemoryBudget::new(1000)?; // 1KB limit
384
385        // First allocation should succeed
386        let alloc1 = budget.try_allocate(500, "reader1")?;
387        assert!(alloc1.is_some());
388
389        // Second allocation within limit should succeed
390        let alloc2 = budget.try_allocate(400, "reader2")?;
391        assert!(alloc2.is_some());
392
393        // Third allocation exceeding limit should fail
394        let alloc3 = budget.try_allocate(200, "reader3")?;
395        assert!(alloc3.is_none());
396
397        // After dropping first allocation, new allocation should succeed
398        drop(alloc1);
399        let alloc4 = budget.try_allocate(300, "reader4")?;
400        assert!(alloc4.is_some());
401
402        Ok(())
403    }
404
405    #[test]
406    fn memory_pressure_calculation_works() -> Result<(), TaleError> {
407        let budget = MemoryBudget::new(1000)?;
408
409        // Low pressure (< 60%)
410        let _alloc1 = budget.try_allocate(500, "reader1")?;
411        assert_eq!(budget.current_pressure()?, MemoryPressure::Low);
412
413        // Moderate pressure (60-85%)
414        let _alloc2 = budget.try_allocate(150, "reader2")?;
415        assert_eq!(budget.current_pressure()?, MemoryPressure::Moderate);
416
417        // High pressure (85-95%)
418        let _alloc3 = budget.try_allocate(200, "reader3")?;
419        assert_eq!(budget.current_pressure()?, MemoryPressure::High);
420
421        // Critical pressure (> 95%)
422        let _alloc4 = budget.try_allocate(100, "reader4")?;
423        assert_eq!(budget.current_pressure()?, MemoryPressure::Critical);
424
425        Ok(())
426    }
427
428    #[test]
429    fn can_recommend_chunk_size() -> Result<(), TaleError> {
430        let budget = MemoryBudget::new(1000)?;
431
432        // Low pressure - no reduction
433        assert_eq!(budget.recommended_chunk_size(1000)?, 1000);
434
435        // Force moderate pressure
436        let _alloc = budget.try_allocate(700, "reader1")?;
437        assert_eq!(budget.recommended_chunk_size(1000)?, 800); // 20% reduction
438
439        Ok(())
440    }
441
442    #[test]
443    fn allocation_automatic_cleanup_works() -> Result<(), TaleError> {
444        let budget = MemoryBudget::new(1000)?;
445
446        {
447            let _alloc1 = budget.try_allocate(500, "reader1")?;
448            let _alloc2 = budget.try_allocate(400, "reader2")?;
449            let stats = budget.usage_stats()?;
450            assert_eq!(stats.current_usage, 900);
451        } // Allocations dropped here
452
453        // After drop, usage should be 0
454        let stats = budget.usage_stats()?;
455        assert_eq!(stats.current_usage, 0);
456
457        Ok(())
458    }
459}
tale_ndjson/memory_budget.rs

tale_ndjson/
memory_budget.rs
