Skip to main content

ferrum_runtime/memory/
stats.rs

1//! Memory statistics and monitoring
2
3use ferrum_types::Device;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
7use std::time::{Duration, Instant};
8
9/// Detailed memory statistics for monitoring
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct DetailedMemoryStats {
12    /// Device this stats belong to
13    pub device: Device,
14    /// Total memory allocated (lifetime)
15    pub total_allocated_bytes: u64,
16    /// Total memory deallocated (lifetime)  
17    pub total_deallocated_bytes: u64,
18    /// Current memory usage
19    pub current_usage_bytes: usize,
20    /// Peak memory usage
21    pub peak_usage_bytes: usize,
22    /// Number of active allocations
23    pub active_allocations: usize,
24    /// Number of allocation requests
25    pub allocation_count: u64,
26    /// Number of deallocation requests
27    pub deallocation_count: u64,
28    /// Number of allocation failures
29    pub allocation_failures: u64,
30    /// Average allocation size
31    pub avg_allocation_size: f64,
32    /// Fragmentation ratio (0.0 - 1.0)
33    pub fragmentation_ratio: f32,
34    /// Time since last reset
35    pub uptime: Duration,
36    /// Allocation size histogram
37    pub size_histogram: HashMap<String, u64>,
38}
39
40/// Memory statistics tracker
41pub struct MemoryStatsTracker {
42    device: Device,
43    total_allocated: AtomicU64,
44    total_deallocated: AtomicU64,
45    current_usage: AtomicUsize,
46    peak_usage: AtomicUsize,
47    active_allocations: AtomicUsize,
48    allocation_count: AtomicU64,
49    deallocation_count: AtomicU64,
50    allocation_failures: AtomicU64,
51    start_time: Instant,
52    size_buckets: parking_lot::Mutex<[u64; 16]>, // Size histogram buckets
53}
54
55impl MemoryStatsTracker {
56    /// Create new memory stats tracker
57    pub fn new(device: Device) -> Self {
58        Self {
59            device,
60            total_allocated: AtomicU64::new(0),
61            total_deallocated: AtomicU64::new(0),
62            current_usage: AtomicUsize::new(0),
63            peak_usage: AtomicUsize::new(0),
64            active_allocations: AtomicUsize::new(0),
65            allocation_count: AtomicU64::new(0),
66            deallocation_count: AtomicU64::new(0),
67            allocation_failures: AtomicU64::new(0),
68            start_time: Instant::now(),
69            size_buckets: parking_lot::Mutex::new([0; 16]),
70        }
71    }
72
73    /// Record allocation
74    pub fn record_allocation(&self, size: usize) {
75        self.total_allocated
76            .fetch_add(size as u64, Ordering::Relaxed);
77        self.allocation_count.fetch_add(1, Ordering::Relaxed);
78        self.active_allocations.fetch_add(1, Ordering::Relaxed);
79
80        let new_usage = self.current_usage.fetch_add(size, Ordering::Relaxed) + size;
81
82        // Update peak if necessary
83        let mut peak = self.peak_usage.load(Ordering::Relaxed);
84        while new_usage > peak {
85            match self.peak_usage.compare_exchange_weak(
86                peak,
87                new_usage,
88                Ordering::Relaxed,
89                Ordering::Relaxed,
90            ) {
91                Ok(_) => break,
92                Err(current_peak) => peak = current_peak,
93            }
94        }
95
96        // Update size histogram
97        self.update_size_histogram(size);
98    }
99
100    /// Record deallocation
101    pub fn record_deallocation(&self, size: usize) {
102        self.total_deallocated
103            .fetch_add(size as u64, Ordering::Relaxed);
104        self.deallocation_count.fetch_add(1, Ordering::Relaxed);
105        self.active_allocations.fetch_sub(1, Ordering::Relaxed);
106        self.current_usage.fetch_sub(size, Ordering::Relaxed);
107    }
108
109    /// Record allocation failure
110    pub fn record_allocation_failure(&self) {
111        self.allocation_failures.fetch_add(1, Ordering::Relaxed);
112    }
113
114    /// Get current statistics
115    pub fn stats(&self) -> DetailedMemoryStats {
116        let total_allocated = self.total_allocated.load(Ordering::Relaxed);
117        let allocation_count = self.allocation_count.load(Ordering::Relaxed);
118
119        let avg_allocation_size = if allocation_count > 0 {
120            total_allocated as f64 / allocation_count as f64
121        } else {
122            0.0
123        };
124
125        // Build size histogram
126        let buckets = self.size_buckets.lock();
127        let mut size_histogram = HashMap::new();
128        let bucket_labels = [
129            "0-1KB",
130            "1KB-4KB",
131            "4KB-16KB",
132            "16KB-64KB",
133            "64KB-256KB",
134            "256KB-1MB",
135            "1MB-4MB",
136            "4MB-16MB",
137            "16MB-64MB",
138            "64MB-256MB",
139            "256MB-1GB",
140            "1GB-4GB",
141            "4GB-16GB",
142            "16GB-64GB",
143            "64GB+",
144            "Other",
145        ];
146
147        for (i, &count) in buckets.iter().enumerate() {
148            if count > 0 {
149                size_histogram.insert(bucket_labels[i].to_string(), count);
150            }
151        }
152
153        DetailedMemoryStats {
154            device: self.device.clone(),
155            total_allocated_bytes: total_allocated,
156            total_deallocated_bytes: self.total_deallocated.load(Ordering::Relaxed),
157            current_usage_bytes: self.current_usage.load(Ordering::Relaxed),
158            peak_usage_bytes: self.peak_usage.load(Ordering::Relaxed),
159            active_allocations: self.active_allocations.load(Ordering::Relaxed),
160            allocation_count,
161            deallocation_count: self.deallocation_count.load(Ordering::Relaxed),
162            allocation_failures: self.allocation_failures.load(Ordering::Relaxed),
163            avg_allocation_size,
164            fragmentation_ratio: self.calculate_fragmentation_ratio(),
165            uptime: self.start_time.elapsed(),
166            size_histogram,
167        }
168    }
169
170    /// Reset statistics
171    pub fn reset(&self) {
172        self.total_allocated.store(0, Ordering::Relaxed);
173        self.total_deallocated.store(0, Ordering::Relaxed);
174        self.allocation_count.store(0, Ordering::Relaxed);
175        self.deallocation_count.store(0, Ordering::Relaxed);
176        self.allocation_failures.store(0, Ordering::Relaxed);
177
178        let mut buckets = self.size_buckets.lock();
179        buckets.fill(0);
180    }
181
182    fn update_size_histogram(&self, size: usize) {
183        let bucket_index = match size {
184            0..=1024 => 0,                   // 0-1KB
185            1025..=4096 => 1,                // 1KB-4KB
186            4097..=16384 => 2,               // 4KB-16KB
187            16385..=65536 => 3,              // 16KB-64KB
188            65537..=262144 => 4,             // 64KB-256KB
189            262145..=1048576 => 5,           // 256KB-1MB
190            1048577..=4194304 => 6,          // 1MB-4MB
191            4194305..=16777216 => 7,         // 4MB-16MB
192            16777217..=67108864 => 8,        // 16MB-64MB
193            67108865..=268435456 => 9,       // 64MB-256MB
194            268435457..=1073741824 => 10,    // 256MB-1GB
195            1073741825..=4294967296 => 11,   // 1GB-4GB
196            4294967297..=17179869184 => 12,  // 4GB-16GB
197            17179869185..=68719476736 => 13, // 16GB-64GB
198            68719476737.. => 14,             // 64GB+
199        };
200
201        let mut buckets = self.size_buckets.lock();
202        buckets[bucket_index] += 1;
203    }
204
205    fn calculate_fragmentation_ratio(&self) -> f32 {
206        // Simplified fragmentation calculation
207        // In a real implementation, this would analyze actual memory layout
208        let total_allocated = self.total_allocated.load(Ordering::Relaxed);
209        let total_deallocated = self.total_deallocated.load(Ordering::Relaxed);
210        let current_usage = self.current_usage.load(Ordering::Relaxed);
211
212        if total_allocated == 0 {
213            return 0.0;
214        }
215
216        // Fragmentation increases with the ratio of deallocated to allocated memory
217        let dealloc_ratio = total_deallocated as f32 / total_allocated as f32;
218        let usage_ratio = current_usage as f32 / total_allocated as f32;
219
220        // Simple heuristic: fragmentation is higher when we've deallocated a lot
221        // but still have high current usage
222        (dealloc_ratio * (1.0 - usage_ratio)).min(1.0)
223    }
224}
225
226/// Global memory statistics registry
227pub struct GlobalMemoryStatsRegistry {
228    trackers: parking_lot::RwLock<HashMap<Device, MemoryStatsTracker>>,
229}
230
231impl GlobalMemoryStatsRegistry {
232    /// Create new global registry
233    pub fn new() -> Self {
234        Self {
235            trackers: parking_lot::RwLock::new(HashMap::new()),
236        }
237    }
238
239    /// Get or create tracker for device
240    pub fn get_or_create_tracker(
241        &self,
242        device: Device,
243    ) -> parking_lot::RwLockReadGuard<'_, MemoryStatsTracker> {
244        // First try to get existing tracker
245        {
246            let trackers = self.trackers.read();
247            if trackers.contains_key(&device) {
248                // Can't return the reference directly due to lifetime issues
249                // This is a simplified implementation
250            }
251        }
252
253        // Create new tracker if needed
254        {
255            let mut trackers = self.trackers.write();
256            trackers
257                .entry(device.clone())
258                .or_insert_with(|| MemoryStatsTracker::new(device));
259        }
260
261        // Return read guard (simplified)
262        unimplemented!("Simplified implementation - would need proper lifetime management")
263    }
264
265    /// Get stats for device
266    pub fn get_stats(&self, device: Device) -> Option<DetailedMemoryStats> {
267        let trackers = self.trackers.read();
268        trackers.get(&device).map(|t| t.stats())
269    }
270
271    /// Get stats for all devices
272    pub fn get_all_stats(&self) -> Vec<DetailedMemoryStats> {
273        let trackers = self.trackers.read();
274        trackers.values().map(|t| t.stats()).collect()
275    }
276
277    /// Reset stats for device
278    pub fn reset_stats(&self, device: Device) {
279        let trackers = self.trackers.read();
280        if let Some(tracker) = trackers.get(&device) {
281            tracker.reset();
282        }
283    }
284
285    /// Reset all stats
286    pub fn reset_all_stats(&self) {
287        let trackers = self.trackers.read();
288        for tracker in trackers.values() {
289            tracker.reset();
290        }
291    }
292}
293
294impl Default for GlobalMemoryStatsRegistry {
295    fn default() -> Self {
296        Self::new()
297    }
298}
299
300/// Global instance of memory stats registry
301use once_cell::sync::Lazy;
302static GLOBAL_MEMORY_STATS: Lazy<GlobalMemoryStatsRegistry> =
303    Lazy::new(|| GlobalMemoryStatsRegistry::new());
304
305/// Get global memory stats registry
306pub fn global_memory_stats() -> &'static GlobalMemoryStatsRegistry {
307    &GLOBAL_MEMORY_STATS
308}
309
310#[cfg(test)]
311mod tests {
312    use super::*;
313
314    #[test]
315    fn test_memory_stats_tracker_creation() {
316        let tracker = MemoryStatsTracker::new(Device::CPU);
317        let stats = tracker.stats();
318
319        assert_eq!(stats.total_allocated_bytes, 0);
320        assert_eq!(stats.current_usage_bytes, 0);
321        assert_eq!(stats.active_allocations, 0);
322        assert_eq!(stats.allocation_count, 0);
323    }
324
325    #[test]
326    fn test_record_allocation() {
327        let tracker = MemoryStatsTracker::new(Device::CPU);
328
329        tracker.record_allocation(1024);
330        let stats = tracker.stats();
331
332        assert_eq!(stats.total_allocated_bytes, 1024);
333        assert_eq!(stats.current_usage_bytes, 1024);
334        assert_eq!(stats.active_allocations, 1);
335        assert_eq!(stats.allocation_count, 1);
336        assert_eq!(stats.peak_usage_bytes, 1024);
337    }
338
339    #[test]
340    fn test_record_deallocation() {
341        let tracker = MemoryStatsTracker::new(Device::CPU);
342
343        tracker.record_allocation(1024);
344        tracker.record_deallocation(1024);
345
346        let stats = tracker.stats();
347        assert_eq!(stats.total_allocated_bytes, 1024);
348        assert_eq!(stats.total_deallocated_bytes, 1024);
349        assert_eq!(stats.current_usage_bytes, 0);
350        assert_eq!(stats.active_allocations, 0);
351        assert_eq!(stats.peak_usage_bytes, 1024); // Peak should remain
352    }
353
354    #[test]
355    fn test_peak_usage_tracking() {
356        let tracker = MemoryStatsTracker::new(Device::CPU);
357
358        tracker.record_allocation(1024);
359        tracker.record_allocation(2048);
360        tracker.record_deallocation(1024);
361        tracker.record_allocation(512);
362
363        let stats = tracker.stats();
364        assert_eq!(stats.peak_usage_bytes, 3072); // 1024 + 2048
365        assert_eq!(stats.current_usage_bytes, 2560); // 2048 + 512
366    }
367
368    #[test]
369    fn test_average_allocation_size() {
370        let tracker = MemoryStatsTracker::new(Device::CPU);
371
372        tracker.record_allocation(1000);
373        tracker.record_allocation(2000);
374        tracker.record_allocation(3000);
375
376        let stats = tracker.stats();
377        assert_eq!(stats.avg_allocation_size, 2000.0);
378    }
379
380    #[test]
381    fn test_allocation_failure_tracking() {
382        let tracker = MemoryStatsTracker::new(Device::CPU);
383
384        tracker.record_allocation_failure();
385        tracker.record_allocation_failure();
386
387        let stats = tracker.stats();
388        assert_eq!(stats.allocation_failures, 2);
389    }
390
391    #[test]
392    fn test_size_histogram() {
393        let tracker = MemoryStatsTracker::new(Device::CPU);
394
395        tracker.record_allocation(512); // 0-1KB
396        tracker.record_allocation(2048); // 1KB-4KB
397        tracker.record_allocation(8192); // 4KB-16KB
398        tracker.record_allocation(32768); // 16KB-64KB
399        tracker.record_allocation(1048576); // 256KB-1MB
400
401        let stats = tracker.stats();
402        assert!(!stats.size_histogram.is_empty());
403        assert!(stats.size_histogram.contains_key("0-1KB"));
404        assert!(stats.size_histogram.contains_key("1KB-4KB"));
405    }
406
407    #[test]
408    fn test_reset_stats() {
409        let tracker = MemoryStatsTracker::new(Device::CPU);
410
411        tracker.record_allocation(1024);
412        tracker.record_allocation(2048);
413        tracker.record_allocation_failure();
414
415        let stats_before = tracker.stats();
416        assert_eq!(stats_before.allocation_count, 2);
417
418        tracker.reset();
419
420        let stats_after = tracker.stats();
421        assert_eq!(stats_after.total_allocated_bytes, 0);
422        assert_eq!(stats_after.allocation_count, 0);
423        assert_eq!(stats_after.allocation_failures, 0);
424    }
425
426    #[test]
427    fn test_fragmentation_ratio() {
428        let tracker = MemoryStatsTracker::new(Device::CPU);
429
430        // Initially no fragmentation
431        let stats = tracker.stats();
432        assert_eq!(stats.fragmentation_ratio, 0.0);
433
434        // Allocate and deallocate to create fragmentation
435        tracker.record_allocation(1000);
436        tracker.record_allocation(2000);
437        tracker.record_deallocation(1000);
438
439        let stats = tracker.stats();
440        assert!(stats.fragmentation_ratio > 0.0);
441        assert!(stats.fragmentation_ratio <= 1.0);
442    }
443
444    #[test]
445    fn test_global_memory_stats_registry() {
446        let registry = GlobalMemoryStatsRegistry::new();
447
448        // Initially empty
449        let stats = registry.get_stats(Device::CPU);
450        assert!(stats.is_none());
451
452        // Get all stats (should be empty)
453        let all_stats = registry.get_all_stats();
454        assert_eq!(all_stats.len(), 0);
455    }
456
457    #[test]
458    fn test_global_memory_stats_singleton() {
459        let registry1 = global_memory_stats();
460        let registry2 = global_memory_stats();
461
462        // Should be the same instance
463        assert!(std::ptr::eq(registry1, registry2));
464    }
465
466    #[test]
467    fn test_uptime_tracking() {
468        let tracker = MemoryStatsTracker::new(Device::CPU);
469
470        // Sleep briefly to ensure uptime is non-zero
471        std::thread::sleep(std::time::Duration::from_millis(10));
472
473        let stats = tracker.stats();
474        assert!(stats.uptime.as_millis() >= 10);
475    }
476
477    #[test]
478    fn test_multiple_allocations_deallocations() {
479        let tracker = MemoryStatsTracker::new(Device::CPU);
480
481        // Simulate multiple allocation/deallocation cycles
482        for i in 1..=10 {
483            tracker.record_allocation(i * 100);
484        }
485
486        let stats_after_alloc = tracker.stats();
487        assert_eq!(stats_after_alloc.allocation_count, 10);
488        assert_eq!(stats_after_alloc.active_allocations, 10);
489
490        for _ in 1..=5 {
491            tracker.record_deallocation(100);
492        }
493
494        let stats_after_dealloc = tracker.stats();
495        assert_eq!(stats_after_dealloc.deallocation_count, 5);
496        assert_eq!(stats_after_dealloc.active_allocations, 5);
497    }
498}