memscope_rs/snapshot/
engine.rs

1//! Snapshot Engine - Snapshot construction and aggregation
2//!
3//! This module provides the SnapshotEngine which is responsible for
4//! building memory snapshots from event data.
5
6use crate::event_store::{MemoryEvent, MemoryEventType, SharedEventStore};
7use crate::snapshot::types::{ActiveAllocation, MemorySnapshot, ThreadMemoryStats};
8use std::collections::HashMap;
9
10/// Snapshot Engine - Builds memory snapshots from event data
11///
12/// The SnapshotEngine is responsible for constructing point-in-time
13/// views of memory usage from the events stored in the EventStore.
14///
15/// Key properties:
16/// - Read-only: Does not consume events from EventStore
17/// - Efficient: Optimized for fast snapshot construction
18/// - Comprehensive: Captures all relevant memory state
19pub struct SnapshotEngine {
20    /// Reference to the event store
21    event_store: SharedEventStore,
22}
23
24impl SnapshotEngine {
25    /// Create a new SnapshotEngine
26    pub fn new(event_store: SharedEventStore) -> Self {
27        Self { event_store }
28    }
29
30    /// Build a snapshot from the current event store state
31    ///
32    /// This method reads all events from the event store and
33    /// constructs a snapshot representing the current memory state.
34    pub fn build_snapshot(&self) -> MemorySnapshot {
35        let events = self.event_store.snapshot();
36        self.build_snapshot_from_events(events)
37    }
38
39    /// Build a snapshot from a specific set of events
40    ///
41    /// # Arguments
42    /// * `events` - The events to build the snapshot from
43    pub fn build_snapshot_from_events(&self, events: Vec<MemoryEvent>) -> MemorySnapshot {
44        let mut snapshot = MemorySnapshot::new();
45        let mut ptr_to_allocation: HashMap<usize, ActiveAllocation> = HashMap::new();
46        let mut thread_stats: HashMap<u64, ThreadMemoryStats> = HashMap::new();
47        let mut peak_memory: usize = 0;
48        let mut current_memory: usize = 0;
49
50        for event in events {
51            match event.event_type {
52                MemoryEventType::Allocate => {
53                    let allocation = ActiveAllocation {
54                        ptr: event.ptr,
55                        size: event.size,
56                        allocated_at: event.timestamp,
57                        var_name: event.var_name,
58                        type_name: event.type_name,
59                        thread_id: event.thread_id,
60                        call_stack_hash: event.call_stack_hash,
61                    };
62
63                    ptr_to_allocation.insert(event.ptr, allocation);
64
65                    snapshot.stats.total_allocations += 1;
66                    snapshot.stats.total_allocated += event.size;
67                    current_memory += event.size;
68
69                    let thread_stat =
70                        thread_stats
71                            .entry(event.thread_id)
72                            .or_insert_with(|| ThreadMemoryStats {
73                                thread_id: event.thread_id,
74                                allocation_count: 0,
75                                total_allocated: 0,
76                                total_deallocated: 0,
77                                current_memory: 0,
78                                peak_memory: 0,
79                            });
80                    thread_stat.allocation_count += 1;
81                    thread_stat.total_allocated += event.size;
82                    thread_stat.current_memory += event.size;
83                    if thread_stat.current_memory > thread_stat.peak_memory {
84                        thread_stat.peak_memory = thread_stat.current_memory;
85                    }
86                }
87                MemoryEventType::Reallocate => {
88                    let old_allocation = ptr_to_allocation.get(&event.ptr).cloned();
89                    let old_size = event.old_size.unwrap_or_else(|| {
90                        old_allocation.as_ref().map(|a| a.size).unwrap_or_else(|| {
91                            tracing::warn!(
92                                "Reallocation without old_size or previous allocation: ptr={:#x}, new_size={}",
93                                event.ptr,
94                                event.size
95                            );
96                            0
97                        })
98                    });
99
100                    let allocation = ActiveAllocation {
101                        ptr: event.ptr,
102                        size: event.size,
103                        allocated_at: old_allocation
104                            .map(|a| a.allocated_at)
105                            .unwrap_or(event.timestamp),
106                        var_name: event.var_name,
107                        type_name: event.type_name,
108                        thread_id: event.thread_id,
109                        call_stack_hash: event.call_stack_hash,
110                    };
111
112                    ptr_to_allocation.insert(event.ptr, allocation);
113
114                    snapshot.stats.total_reallocations += 1;
115                    snapshot.stats.total_allocated += event.size;
116                    snapshot.stats.total_deallocated += old_size;
117
118                    // Validate memory calculation to detect potential bugs
119                    if old_size > current_memory {
120                        tracing::warn!(
121                            "Potential memory tracking bug detected: deallocating {} bytes but only {} bytes are currently tracked. ptr=0x{:x}",
122                            old_size, current_memory, event.ptr
123                        );
124                        // This could indicate a bug (double-free, incorrect tracking, etc.)
125                        // but we still want to update the calculation safely
126                        current_memory = current_memory
127                            .saturating_sub(old_size)
128                            .saturating_add(event.size);
129                    } else {
130                        // Normal case: safely perform the calculation
131                        current_memory =
132                            current_memory
133                                .checked_sub(old_size)
134                                .and_then(|v| v.checked_add(event.size))
135                                .unwrap_or_else(|| {
136                                    tracing::error!(
137                                    "Integer overflow detected in memory calculation: {} - {} + {}",
138                                    current_memory, old_size, event.size
139                                );
140                                    // Fallback to saturating operations to prevent panic
141                                    current_memory
142                                        .saturating_sub(old_size)
143                                        .saturating_add(event.size)
144                                });
145                    }
146
147                    let thread_stat =
148                        thread_stats
149                            .entry(event.thread_id)
150                            .or_insert_with(|| ThreadMemoryStats {
151                                thread_id: event.thread_id,
152                                allocation_count: 0,
153                                total_allocated: 0,
154                                total_deallocated: 0,
155                                current_memory: 0,
156                                peak_memory: 0,
157                            });
158                    thread_stat.total_allocated += event.size;
159                    thread_stat.total_deallocated += old_size;
160                    thread_stat.current_memory = thread_stat
161                        .current_memory
162                        .saturating_sub(old_size)
163                        .saturating_add(event.size);
164                    if thread_stat.current_memory > thread_stat.peak_memory {
165                        thread_stat.peak_memory = thread_stat.current_memory;
166                    }
167                }
168                MemoryEventType::Deallocate => {
169                    if let Some(allocation) = ptr_to_allocation.remove(&event.ptr) {
170                        snapshot.stats.total_deallocations += 1;
171                        snapshot.stats.total_deallocated += allocation.size;
172                        current_memory = current_memory.saturating_sub(allocation.size);
173
174                        if let Some(thread_stat) = thread_stats.get_mut(&event.thread_id) {
175                            thread_stat.total_deallocated += allocation.size;
176                            thread_stat.current_memory =
177                                thread_stat.current_memory.saturating_sub(allocation.size);
178                        }
179                    } else {
180                        snapshot.stats.unmatched_deallocations += 1;
181                        tracing::debug!(
182                            "Unmatched deallocation: ptr={:#x}, thread_id={}",
183                            event.ptr,
184                            event.thread_id
185                        );
186                    }
187                }
188                MemoryEventType::Move | MemoryEventType::Borrow | MemoryEventType::Return => {
189                    // These don't affect the current memory state
190                    // but we may want to track them for analysis
191                }
192            }
193
194            // Update peak memory
195            if current_memory > peak_memory {
196                peak_memory = current_memory;
197            }
198        }
199
200        // Build final snapshot
201        snapshot.active_allocations = ptr_to_allocation;
202        snapshot.thread_stats = thread_stats;
203        snapshot.stats.active_allocations = snapshot.active_allocations.len();
204        snapshot.stats.current_memory = current_memory;
205        snapshot.stats.peak_memory = peak_memory;
206
207        snapshot
208    }
209
210    /// Get the event store reference
211    pub fn event_store(&self) -> &SharedEventStore {
212        &self.event_store
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use crate::event_store::EventStore;
220    use std::sync::Arc;
221
222    #[test]
223    fn test_snapshot_engine_creation() {
224        let event_store = Arc::new(EventStore::new());
225        let engine = SnapshotEngine::new(event_store);
226        let snapshot = engine.build_snapshot();
227        assert_eq!(snapshot.active_count(), 0);
228    }
229
230    #[test]
231    fn test_snapshot_with_allocations() {
232        let event_store = Arc::new(EventStore::new());
233        event_store.record(MemoryEvent::allocate(0x1000, 1024, 1));
234        event_store.record(MemoryEvent::allocate(0x2000, 2048, 1));
235
236        let engine = SnapshotEngine::new(event_store);
237        let snapshot = engine.build_snapshot();
238
239        assert_eq!(snapshot.active_count(), 2);
240        assert_eq!(snapshot.current_memory(), 3072);
241    }
242
243    #[test]
244    fn test_snapshot_with_deallocations() {
245        let event_store = Arc::new(EventStore::new());
246        event_store.record(MemoryEvent::allocate(0x1000, 1024, 1));
247        event_store.record(MemoryEvent::deallocate(0x1000, 1024, 1));
248
249        let engine = SnapshotEngine::new(event_store);
250        let snapshot = engine.build_snapshot();
251
252        assert_eq!(snapshot.active_count(), 0);
253        assert_eq!(snapshot.current_memory(), 0);
254        assert_eq!(snapshot.stats.total_allocations, 1);
255        assert_eq!(snapshot.stats.total_deallocations, 1);
256    }
257
258    #[test]
259    fn test_snapshot_peak_memory() {
260        let event_store = Arc::new(EventStore::new());
261        event_store.record(MemoryEvent::allocate(0x1000, 1024, 1));
262        event_store.record(MemoryEvent::allocate(0x2000, 2048, 1));
263        event_store.record(MemoryEvent::deallocate(0x2000, 2048, 1));
264
265        let engine = SnapshotEngine::new(event_store);
266        let snapshot = engine.build_snapshot();
267
268        assert_eq!(snapshot.peak_memory(), 3072);
269        assert_eq!(snapshot.current_memory(), 1024);
270    }
271
272    #[test]
273    fn test_snapshot_thread_stats() {
274        let event_store = Arc::new(EventStore::new());
275        event_store.record(MemoryEvent::allocate(0x1000, 1024, 1));
276        event_store.record(MemoryEvent::allocate(0x2000, 2048, 2));
277
278        let engine = SnapshotEngine::new(event_store);
279        let snapshot = engine.build_snapshot();
280
281        assert_eq!(snapshot.thread_stats.len(), 2);
282
283        let thread1 = snapshot.thread_stats.get(&1).unwrap();
284        assert_eq!(thread1.allocation_count, 1);
285        assert_eq!(thread1.total_allocated, 1024);
286
287        let thread2 = snapshot.thread_stats.get(&2).unwrap();
288        assert_eq!(thread2.allocation_count, 1);
289        assert_eq!(thread2.total_allocated, 2048);
290    }
291
292    #[test]
293    fn test_deallocation_underflow_protection() {
294        let event_store = Arc::new(EventStore::new());
295        event_store.record(MemoryEvent::allocate(0x1000, 1024, 1));
296        event_store.record(MemoryEvent::deallocate(0x1000, 2048, 1));
297
298        let engine = SnapshotEngine::new(event_store);
299        let snapshot = engine.build_snapshot();
300
301        assert_eq!(snapshot.current_memory(), 0);
302        let thread1 = snapshot.thread_stats.get(&1).unwrap();
303        assert_eq!(thread1.current_memory, 0);
304    }
305}
memscope_rs/snapshot/engine.rs

memscope_rs/snapshot/
engine.rs
