1use 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#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct DetailedMemoryStats {
12 pub device: Device,
14 pub total_allocated_bytes: u64,
16 pub total_deallocated_bytes: u64,
18 pub current_usage_bytes: usize,
20 pub peak_usage_bytes: usize,
22 pub active_allocations: usize,
24 pub allocation_count: u64,
26 pub deallocation_count: u64,
28 pub allocation_failures: u64,
30 pub avg_allocation_size: f64,
32 pub fragmentation_ratio: f32,
34 pub uptime: Duration,
36 pub size_histogram: HashMap<String, u64>,
38}
39
40pub 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]>, }
54
55impl MemoryStatsTracker {
56 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 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 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 self.update_size_histogram(size);
98 }
99
100 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 pub fn record_allocation_failure(&self) {
111 self.allocation_failures.fetch_add(1, Ordering::Relaxed);
112 }
113
114 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 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 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, 1025..=4096 => 1, 4097..=16384 => 2, 16385..=65536 => 3, 65537..=262144 => 4, 262145..=1048576 => 5, 1048577..=4194304 => 6, 4194305..=16777216 => 7, 16777217..=67108864 => 8, 67108865..=268435456 => 9, 268435457..=1073741824 => 10, 1073741825..=4294967296 => 11, 4294967297..=17179869184 => 12, 17179869185..=68719476736 => 13, 68719476737.. => 14, };
200
201 let mut buckets = self.size_buckets.lock();
202 buckets[bucket_index] += 1;
203 }
204
205 fn calculate_fragmentation_ratio(&self) -> f32 {
206 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 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 (dealloc_ratio * (1.0 - usage_ratio)).min(1.0)
223 }
224}
225
226pub struct GlobalMemoryStatsRegistry {
228 trackers: parking_lot::RwLock<HashMap<Device, MemoryStatsTracker>>,
229}
230
231impl GlobalMemoryStatsRegistry {
232 pub fn new() -> Self {
234 Self {
235 trackers: parking_lot::RwLock::new(HashMap::new()),
236 }
237 }
238
239 pub fn get_or_create_tracker(
241 &self,
242 device: Device,
243 ) -> parking_lot::RwLockReadGuard<'_, MemoryStatsTracker> {
244 {
246 let trackers = self.trackers.read();
247 if trackers.contains_key(&device) {
248 }
251 }
252
253 {
255 let mut trackers = self.trackers.write();
256 trackers
257 .entry(device.clone())
258 .or_insert_with(|| MemoryStatsTracker::new(device));
259 }
260
261 unimplemented!("Simplified implementation - would need proper lifetime management")
263 }
264
265 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 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 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 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
300use once_cell::sync::Lazy;
302static GLOBAL_MEMORY_STATS: Lazy<GlobalMemoryStatsRegistry> =
303 Lazy::new(GlobalMemoryStatsRegistry::new);
304
305pub 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); }
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); assert_eq!(stats.current_usage_bytes, 2560); }
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); tracker.record_allocation(2048); tracker.record_allocation(8192); tracker.record_allocation(32768); tracker.record_allocation(1048576); 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 let stats = tracker.stats();
432 assert_eq!(stats.fragmentation_ratio, 0.0);
433
434 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 let stats = registry.get_stats(Device::CPU);
450 assert!(stats.is_none());
451
452 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 assert!(std::ptr::eq(registry1, registry2));
464 }
465
466 #[test]
467 fn test_uptime_tracking() {
468 let tracker = MemoryStatsTracker::new(Device::CPU);
469
470 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 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}