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