Skip to main content

memscope_rs/snapshot/
types.rs

1//! Snapshot types - Data structures for memory snapshots
2//!
3//! This module defines the core data structures used by the
4//! SnapshotEngine for representing memory snapshots.
5
6use crate::core::types::TrackKind;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10/// Active allocation information in a snapshot
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
12pub struct ActiveAllocation {
13    /// Memory pointer address (None for Container/Value types)
14    pub ptr: Option<usize>,
15    /// Allocation size in bytes
16    pub size: usize,
17    /// Memory allocation semantic role
18    pub kind: TrackKind,
19    /// Timestamp when this allocation was made
20    pub allocated_at: u64,
21    /// Optional variable name
22    pub var_name: Option<String>,
23    /// Optional type name
24    pub type_name: Option<String>,
25    /// Thread ID that made this allocation
26    pub thread_id: u64,
27    /// Optional call stack hash for clone detection
28    pub call_stack_hash: Option<u64>,
29    /// Module path (from module_path!())
30    pub module_path: Option<String>,
31    /// Stack pointer (for StackOwner types like Arc/Rc)
32    pub stack_ptr: Option<usize>,
33}
34
35/// Memory statistics for a snapshot
36#[derive(Debug, Clone, Serialize, Deserialize, Default)]
37pub struct MemoryStats {
38    /// Total number of allocations in the snapshot
39    pub total_allocations: usize,
40    /// Total number of reallocations
41    pub total_reallocations: usize,
42    /// Total number of deallocations
43    pub total_deallocations: usize,
44    /// Number of deallocations without matching allocations
45    pub unmatched_deallocations: usize,
46    /// Number of currently active allocations
47    pub active_allocations: usize,
48    /// Total bytes allocated
49    pub total_allocated: usize,
50    /// Total bytes deallocated
51    pub total_deallocated: usize,
52    /// Currently used memory (sum of active allocations)
53    pub current_memory: usize,
54    /// Peak memory usage observed
55    pub peak_memory: usize,
56}
57
58/// Thread-specific memory statistics
59#[derive(Debug, Clone, Serialize, Deserialize, Default)]
60pub struct ThreadMemoryStats {
61    /// Thread ID
62    pub thread_id: u64,
63    /// Number of allocations by this thread
64    pub allocation_count: usize,
65    /// Total bytes allocated by this thread
66    pub total_allocated: usize,
67    /// Total bytes deallocated by this thread
68    pub total_deallocated: usize,
69    /// Current memory usage by this thread
70    pub current_memory: usize,
71    /// Peak memory usage by this thread
72    pub peak_memory: usize,
73}
74
75/// Memory snapshot - a point-in-time view of memory usage
76#[derive(Debug, Clone, Serialize, Deserialize, Default)]
77pub struct MemorySnapshot {
78    /// Timestamp when this snapshot was taken
79    pub timestamp: u64,
80    /// Overall memory statistics
81    pub stats: MemoryStats,
82    /// Active allocations (ptr -> allocation info)
83    pub active_allocations: HashMap<usize, ActiveAllocation>,
84    /// Per-thread statistics
85    pub thread_stats: HashMap<u64, ThreadMemoryStats>,
86}
87
88impl MemorySnapshot {
89    /// Create a new empty snapshot
90    pub fn new() -> Self {
91        Self {
92            timestamp: std::time::SystemTime::now()
93                .duration_since(std::time::UNIX_EPOCH)
94                .unwrap_or_default()
95                .as_nanos() as u64,
96            stats: MemoryStats::default(),
97            active_allocations: HashMap::new(),
98            thread_stats: HashMap::new(),
99        }
100    }
101
102    /// Build a MemorySnapshot from a list of AllocationInfo (capture module type)
103    pub fn from_allocation_infos(allocations: Vec<crate::capture::types::AllocationInfo>) -> Self {
104        let mut snapshot = Self::new();
105        let mut thread_stats: HashMap<u64, ThreadMemoryStats> = HashMap::new();
106        let mut current_memory: usize = 0;
107
108        for alloc in allocations {
109            let thread_id = alloc.thread_id_u64;
110
111            let active_alloc = ActiveAllocation {
112                ptr: Some(alloc.ptr),
113                size: alloc.size,
114                kind: TrackKind::HeapOwner {
115                    ptr: alloc.ptr,
116                    size: alloc.size,
117                },
118                allocated_at: alloc.timestamp_alloc,
119                var_name: alloc.var_name,
120                type_name: alloc.type_name,
121                thread_id,
122                call_stack_hash: None,
123                module_path: alloc.module_path,
124                stack_ptr: None,
125            };
126
127            current_memory += alloc.size;
128
129            snapshot.stats.total_allocations += 1;
130            snapshot.stats.total_allocated += alloc.size;
131
132            let thread_stat = thread_stats
133                .entry(thread_id)
134                .or_insert_with(|| ThreadMemoryStats {
135                    thread_id,
136                    allocation_count: 0,
137                    total_allocated: 0,
138                    total_deallocated: 0,
139                    current_memory: 0,
140                    peak_memory: 0,
141                });
142
143            thread_stat.allocation_count += 1;
144            thread_stat.total_allocated += alloc.size;
145            thread_stat.current_memory += alloc.size;
146            if thread_stat.current_memory > thread_stat.peak_memory {
147                thread_stat.peak_memory = thread_stat.current_memory;
148            }
149
150            snapshot.active_allocations.insert(alloc.ptr, active_alloc);
151        }
152
153        snapshot.stats.current_memory = current_memory;
154        snapshot.stats.peak_memory = 0; // Cannot determine peak from current allocations only
155        snapshot.stats.active_allocations = snapshot.active_allocations.len();
156        snapshot.thread_stats = thread_stats;
157
158        snapshot
159    }
160
161    /// Get the number of active allocations
162    pub fn active_count(&self) -> usize {
163        self.active_allocations.len()
164    }
165
166    /// Get the current memory usage
167    pub fn current_memory(&self) -> usize {
168        self.stats.current_memory
169    }
170
171    /// Get the peak memory usage
172    pub fn peak_memory(&self) -> usize {
173        self.stats.peak_memory
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    #[test]
182    fn test_memory_snapshot_new() {
183        let snapshot = MemorySnapshot::new();
184        assert!(snapshot.timestamp > 0);
185        assert_eq!(snapshot.stats.total_allocations, 0);
186        assert!(snapshot.active_allocations.is_empty());
187    }
188
189    #[test]
190    fn test_memory_snapshot_default() {
191        let snapshot = MemorySnapshot::default();
192        assert_eq!(snapshot.timestamp, 0);
193        assert_eq!(snapshot.stats.total_allocations, 0);
194    }
195
196    #[test]
197    fn test_active_allocation_creation() {
198        let alloc = ActiveAllocation {
199            ptr: Some(0x1000),
200            size: 1024,
201            kind: TrackKind::HeapOwner {
202                ptr: 0x1000,
203                size: 1024,
204            },
205            allocated_at: 1000,
206            var_name: Some("test".to_string()),
207            type_name: Some("Vec<u8>".to_string()),
208            thread_id: 1,
209            call_stack_hash: None,
210            module_path: None,
211            stack_ptr: None,
212        };
213
214        assert_eq!(alloc.ptr, Some(0x1000));
215        assert_eq!(alloc.size, 1024);
216        assert_eq!(alloc.var_name, Some("test".to_string()));
217    }
218
219    #[test]
220    fn test_active_allocation_clone() {
221        let alloc = ActiveAllocation {
222            ptr: Some(0x1000),
223            size: 1024,
224            kind: TrackKind::HeapOwner {
225                ptr: 0x1000,
226                size: 1024,
227            },
228            allocated_at: 1000,
229            var_name: None,
230            type_name: None,
231            thread_id: 1,
232            call_stack_hash: None,
233            module_path: None,
234            stack_ptr: None,
235        };
236
237        let cloned = alloc.clone();
238        assert_eq!(cloned.size, alloc.size);
239    }
240
241    #[test]
242    fn test_active_allocation_debug() {
243        let alloc = ActiveAllocation {
244            ptr: Some(0x1000),
245            size: 1024,
246            kind: TrackKind::HeapOwner {
247                ptr: 0x1000,
248                size: 1024,
249            },
250            allocated_at: 1000,
251            var_name: None,
252            type_name: None,
253            thread_id: 1,
254            call_stack_hash: None,
255            module_path: None,
256            stack_ptr: None,
257        };
258
259        let debug_str = format!("{:?}", alloc);
260        assert!(debug_str.contains("ActiveAllocation"));
261    }
262
263    #[test]
264    fn test_memory_stats_default() {
265        let stats = MemoryStats::default();
266        assert_eq!(stats.total_allocations, 0);
267        assert_eq!(stats.total_reallocations, 0);
268        assert_eq!(stats.total_deallocations, 0);
269        assert_eq!(stats.active_allocations, 0);
270        assert_eq!(stats.current_memory, 0);
271        assert_eq!(stats.peak_memory, 0);
272    }
273
274    #[test]
275    fn test_memory_stats_clone() {
276        let stats = MemoryStats {
277            total_allocations: 100,
278            total_reallocations: 10,
279            total_deallocations: 50,
280            unmatched_deallocations: 2,
281            active_allocations: 50,
282            total_allocated: 1024 * 1024,
283            total_deallocated: 512 * 1024,
284            current_memory: 512 * 1024,
285            peak_memory: 1024 * 1024,
286        };
287
288        let cloned = stats.clone();
289        assert_eq!(cloned.total_allocations, 100);
290        assert_eq!(cloned.peak_memory, 1024 * 1024);
291    }
292
293    #[test]
294    fn test_thread_memory_stats_default() {
295        let stats = ThreadMemoryStats::default();
296        assert_eq!(stats.thread_id, 0);
297        assert_eq!(stats.allocation_count, 0);
298        assert_eq!(stats.current_memory, 0);
299    }
300
301    #[test]
302    fn test_thread_memory_stats_clone() {
303        let stats = ThreadMemoryStats {
304            thread_id: 1,
305            allocation_count: 50,
306            total_allocated: 4096,
307            total_deallocated: 2048,
308            current_memory: 2048,
309            peak_memory: 4096,
310        };
311
312        let cloned = stats.clone();
313        assert_eq!(cloned.thread_id, 1);
314        assert_eq!(cloned.allocation_count, 50);
315    }
316
317    #[test]
318    fn test_memory_snapshot_active_count() {
319        let mut snapshot = MemorySnapshot::new();
320        assert_eq!(snapshot.active_count(), 0);
321
322        snapshot.active_allocations.insert(
323            0x1000,
324            ActiveAllocation {
325                ptr: Some(0x1000),
326                size: 1024,
327                kind: TrackKind::HeapOwner {
328                    ptr: 0x1000,
329                    size: 1024,
330                },
331                allocated_at: 1000,
332                var_name: None,
333                type_name: None,
334                thread_id: 1,
335                call_stack_hash: None,
336                module_path: None,
337                stack_ptr: Some(0x2000), // Add stack_ptr field
338            },
339        );
340
341        assert_eq!(snapshot.active_count(), 1);
342    }
343
344    #[test]
345    fn test_memory_snapshot_current_memory() {
346        let mut snapshot = MemorySnapshot::new();
347        assert_eq!(snapshot.current_memory(), 0);
348
349        snapshot.stats.current_memory = 4096;
350        assert_eq!(snapshot.current_memory(), 4096);
351    }
352
353    #[test]
354    fn test_memory_snapshot_peak_memory() {
355        let mut snapshot = MemorySnapshot::new();
356        assert_eq!(snapshot.peak_memory(), 0);
357
358        snapshot.stats.peak_memory = 8192;
359        assert_eq!(snapshot.peak_memory(), 8192);
360    }
361
362    #[test]
363    fn test_memory_snapshot_serialization() {
364        let snapshot = MemorySnapshot::new();
365
366        let json = serde_json::to_string(&snapshot);
367        assert!(json.is_ok());
368
369        let deserialized: Result<MemorySnapshot, _> = serde_json::from_str(&json.unwrap());
370        assert!(deserialized.is_ok());
371    }
372
373    #[test]
374    fn test_active_allocation_serialization() {
375        let alloc = ActiveAllocation {
376            ptr: Some(0x1000),
377            size: 1024,
378            kind: TrackKind::HeapOwner {
379                ptr: 0x1000,
380                size: 1024,
381            },
382            allocated_at: 1000,
383            var_name: Some("test".to_string()),
384            type_name: Some("i32".to_string()),
385            thread_id: 1,
386            call_stack_hash: Some(12345),
387            module_path: Some("test_module".to_string()),
388            stack_ptr: None,
389        };
390
391        let json = serde_json::to_string(&alloc);
392        assert!(json.is_ok());
393
394        let deserialized: Result<ActiveAllocation, _> = serde_json::from_str(&json.unwrap());
395        assert!(deserialized.is_ok());
396    }
397
398    #[test]
399    fn test_memory_stats_serialization() {
400        let stats = MemoryStats {
401            total_allocations: 100,
402            total_reallocations: 10,
403            total_deallocations: 50,
404            unmatched_deallocations: 2,
405            active_allocations: 50,
406            total_allocated: 1024,
407            total_deallocated: 512,
408            current_memory: 512,
409            peak_memory: 1024,
410        };
411
412        let json = serde_json::to_string(&stats);
413        assert!(json.is_ok());
414
415        let deserialized: Result<MemoryStats, _> = serde_json::from_str(&json.unwrap());
416        assert!(deserialized.is_ok());
417    }
418}