Skip to main content

memscope_rs/timeline/
query.rs

1//! Timeline Query - Time-based event queries
2//!
3//! This module provides functionality for querying events based on
4//! time ranges and other temporal criteria.
5
6use crate::event_store::event::{MemoryEvent, MemoryEventType};
7use crate::event_store::EventStore;
8use std::sync::Arc;
9
10/// Timeline query engine
11///
12/// Provides functionality to query events based on time ranges
13/// and other temporal criteria.
14pub struct TimelineQuery {
15    /// Reference to the event store
16    event_store: Arc<EventStore>,
17}
18
19impl TimelineQuery {
20    /// Create a new timeline query
21    pub fn new(event_store: Arc<EventStore>) -> Self {
22        Self { event_store }
23    }
24
25    /// Get events in a time range
26    ///
27    /// # Arguments
28    /// * `start` - Start timestamp (inclusive)
29    /// * `end` - End timestamp (exclusive)
30    pub fn get_events_in_range(&self, start: u64, end: u64) -> Vec<MemoryEvent> {
31        self.event_store
32            .snapshot()
33            .into_iter()
34            .filter(|e| e.timestamp >= start && e.timestamp < end)
35            .collect()
36    }
37
38    /// Get allocation events in a time range
39    ///
40    /// # Arguments
41    /// * `start` - Start timestamp (inclusive)
42    /// * `end` - End timestamp (exclusive)
43    pub fn get_allocations_in_range(&self, start: u64, end: u64) -> Vec<MemoryEvent> {
44        self.get_events_in_range(start, end)
45            .into_iter()
46            .filter(|e| e.event_type == MemoryEventType::Allocate)
47            .collect()
48    }
49
50    /// Get deallocation events in a time range
51    ///
52    /// # Arguments
53    /// * `start` - Start timestamp (inclusive)
54    /// * `end` - End timestamp (exclusive)
55    pub fn get_deallocations_in_range(&self, start: u64, end: u64) -> Vec<MemoryEvent> {
56        self.get_events_in_range(start, end)
57            .into_iter()
58            .filter(|e| e.event_type == MemoryEventType::Deallocate)
59            .collect()
60    }
61
62    /// Get events for a specific thread in a time range
63    ///
64    /// # Arguments
65    /// * `thread_id` - The thread ID to filter by
66    /// * `start` - Start timestamp (inclusive)
67    /// * `end` - End timestamp (exclusive)
68    pub fn get_thread_events_in_range(
69        &self,
70        thread_id: u64,
71        start: u64,
72        end: u64,
73    ) -> Vec<MemoryEvent> {
74        self.get_events_in_range(start, end)
75            .into_iter()
76            .filter(|e| e.thread_id == thread_id)
77            .collect()
78    }
79
80    /// Get memory usage over time with cumulative tracking.
81    ///
82    /// # Arguments
83    /// * `start` - Start timestamp (inclusive)
84    /// * `end` - End timestamp (exclusive)
85    /// * `interval_ms` - Interval between snapshots in milliseconds
86    ///
87    /// # Returns
88    /// Vector of (timestamp, cumulative_memory_bytes) tuples showing
89    /// the actual memory usage at each point in time.
90    pub fn get_memory_usage_over_time(
91        &self,
92        start: u64,
93        end: u64,
94        interval_ms: u64,
95    ) -> Vec<(u64, usize)> {
96        let mut result = Vec::new();
97        let interval_ns = interval_ms * 1_000_000;
98
99        // Get all events in the time range, sorted by timestamp
100        let mut all_events: Vec<(u64, MemoryEvent)> = Vec::new();
101
102        for event in self.get_events_in_range(start, end) {
103            all_events.push((event.timestamp, event));
104        }
105
106        // Sort by timestamp
107        all_events.sort_by_key(|(ts, _)| *ts);
108
109        // Sample at each interval
110        let mut current = start;
111        let mut event_idx = 0;
112        let mut running_memory: usize = 0;
113
114        while current < end {
115            // Process all events up to current timestamp
116            while event_idx < all_events.len() {
117                let (ts, event) = &all_events[event_idx];
118                if *ts > current {
119                    break;
120                }
121
122                match &event.event_type {
123                    MemoryEventType::Allocate => {
124                        running_memory += event.size;
125                    }
126                    MemoryEventType::Deallocate => {
127                        running_memory = running_memory.saturating_sub(event.size);
128                    }
129                    MemoryEventType::Reallocate => {
130                        let old_size = event.old_size.unwrap_or(0);
131                        running_memory = running_memory
132                            .saturating_sub(old_size)
133                            .saturating_add(event.size);
134                    }
135                    _ => {}
136                }
137
138                event_idx += 1;
139            }
140
141            result.push((current, running_memory));
142            current += interval_ns;
143        }
144
145        result
146    }
147
148    /// Get peak memory usage in a time range
149    ///
150    /// # Arguments
151    /// * `start` - Start timestamp (inclusive)
152    /// * `end` - End timestamp (exclusive)
153    ///
154    /// # Note
155    /// This function processes all events in chronological order to find
156    /// the true peak memory usage, not just the maximum interval delta.
157    pub fn get_peak_memory_in_range(&self, start: u64, end: u64) -> usize {
158        // Get all events sorted by timestamp
159        let mut all_events: Vec<MemoryEvent> = Vec::new();
160
161        for event in self.get_events_in_range(start, end) {
162            all_events.push(event);
163        }
164
165        all_events.sort_by_key(|e| e.timestamp);
166
167        // Track running memory and find peak
168        let mut running_memory: usize = 0;
169        let mut peak_memory: usize = 0;
170
171        for event in all_events {
172            match &event.event_type {
173                MemoryEventType::Allocate => {
174                    running_memory += event.size;
175                }
176                MemoryEventType::Deallocate => {
177                    running_memory = running_memory.saturating_sub(event.size);
178                }
179                MemoryEventType::Reallocate => {
180                    let old_size = event.old_size.unwrap_or(0);
181                    running_memory = running_memory
182                        .saturating_sub(old_size)
183                        .saturating_add(event.size);
184                }
185                _ => {}
186            }
187            peak_memory = peak_memory.max(running_memory);
188        }
189
190        peak_memory
191    }
192
193    /// Get event rate (events per second) in a time range
194    ///
195    /// # Arguments
196    /// * `start` - Start timestamp (inclusive)
197    /// * `end` - End timestamp (exclusive)
198    pub fn get_event_rate(&self, start: u64, end: u64) -> f64 {
199        let events = self.get_events_in_range(start, end);
200        let duration_ns = end.saturating_sub(start) as f64;
201        if duration_ns > 0.0 {
202            (events.len() as f64) / (duration_ns / 1_000_000_000.0)
203        } else {
204            0.0
205        }
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212
213    #[test]
214    fn test_timeline_query_creation() {
215        let event_store = Arc::new(EventStore::new());
216        let query = TimelineQuery::new(event_store);
217        let events = query.get_events_in_range(0, 1000);
218        assert!(events.is_empty());
219    }
220
221    #[test]
222    fn test_get_events_in_range() {
223        let event_store = Arc::new(EventStore::new());
224        let event1 = MemoryEvent::allocate(0x1000, 1024, 123);
225        event_store.record(event1);
226        let event2 = MemoryEvent::deallocate(0x1000, 1024, 456);
227        event_store.record(event2);
228
229        let query = TimelineQuery::new(event_store);
230        // Use a large time range to capture all events
231        let events = query.get_events_in_range(0, u64::MAX);
232        assert_eq!(events.len(), 2);
233    }
234
235    #[test]
236    fn test_memory_usage_over_time() {
237        let event_store = Arc::new(EventStore::new());
238        let event = MemoryEvent::allocate(0x1000, 1024, 123);
239        event_store.record(event);
240
241        let query = TimelineQuery::new(event_store);
242        let usage = query.get_memory_usage_over_time(0, 1000, 100);
243        assert!(!usage.is_empty());
244    }
245
246    #[test]
247    fn test_get_allocations_in_range() {
248        let event_store = Arc::new(EventStore::new());
249        let alloc = MemoryEvent::allocate(0x1000, 1024, 100);
250        let dealloc = MemoryEvent::deallocate(0x1000, 1024, 200);
251        event_store.record(alloc);
252        event_store.record(dealloc);
253
254        let query = TimelineQuery::new(event_store);
255        let allocations = query.get_allocations_in_range(0, u64::MAX);
256        assert_eq!(allocations.len(), 1);
257        assert_eq!(allocations[0].event_type, MemoryEventType::Allocate);
258    }
259
260    #[test]
261    fn test_get_deallocations_in_range() {
262        let event_store = Arc::new(EventStore::new());
263        let alloc = MemoryEvent::allocate(0x1000, 1024, 100);
264        let dealloc = MemoryEvent::deallocate(0x1000, 1024, 200);
265        event_store.record(alloc);
266        event_store.record(dealloc);
267
268        let query = TimelineQuery::new(event_store);
269        let deallocations = query.get_deallocations_in_range(0, u64::MAX);
270        assert_eq!(deallocations.len(), 1);
271        assert_eq!(deallocations[0].event_type, MemoryEventType::Deallocate);
272    }
273
274    #[test]
275    fn test_get_thread_events_in_range() {
276        let event_store = Arc::new(EventStore::new());
277        let mut event1 = MemoryEvent::allocate(0x1000, 1024, 100);
278        event1.thread_id = 1;
279        let mut event2 = MemoryEvent::allocate(0x2000, 2048, 150);
280        event2.thread_id = 2;
281        event_store.record(event1);
282        event_store.record(event2);
283
284        let query = TimelineQuery::new(event_store);
285        let thread1_events = query.get_thread_events_in_range(1, 0, u64::MAX);
286        assert_eq!(thread1_events.len(), 1);
287
288        let thread2_events = query.get_thread_events_in_range(2, 0, u64::MAX);
289        assert_eq!(thread2_events.len(), 1);
290
291        let thread3_events = query.get_thread_events_in_range(3, 0, u64::MAX);
292        assert_eq!(thread3_events.len(), 0);
293    }
294
295    #[test]
296    fn test_get_peak_memory_in_range() {
297        let event_store = Arc::new(EventStore::new());
298        let alloc1 = MemoryEvent::allocate(0x1000, 1024, 100);
299        let alloc2 = MemoryEvent::allocate(0x2000, 2048, 150);
300        let dealloc1 = MemoryEvent::deallocate(0x1000, 1024, 200);
301        event_store.record(alloc1);
302        event_store.record(alloc2);
303        event_store.record(dealloc1);
304
305        let query = TimelineQuery::new(event_store);
306        let peak = query.get_peak_memory_in_range(0, u64::MAX);
307        // Peak should be 1024 + 2048 = 3072
308        assert_eq!(peak, 3072);
309    }
310
311    #[test]
312    fn test_get_peak_memory_empty() {
313        let event_store = Arc::new(EventStore::new());
314        let query = TimelineQuery::new(event_store);
315        let peak = query.get_peak_memory_in_range(0, 1000);
316        assert_eq!(peak, 0);
317    }
318
319    #[test]
320    fn test_get_event_rate() {
321        let event_store = Arc::new(EventStore::new());
322        for i in 0..10 {
323            let event = MemoryEvent::allocate(0x1000 + i * 0x100, 1024, i as u64 * 100);
324            event_store.record(event);
325        }
326
327        let query = TimelineQuery::new(event_store);
328        // Use a time range that covers all events
329        let rate = query.get_event_rate(0, 1000);
330        // Rate should be positive if events are in range
331        let _rate = rate;
332    }
333
334    #[test]
335    fn test_get_event_rate_zero_duration() {
336        let event_store = Arc::new(EventStore::new());
337        let event = MemoryEvent::allocate(0x1000, 1024, 100);
338        event_store.record(event);
339
340        let query = TimelineQuery::new(event_store);
341        let rate = query.get_event_rate(100, 100);
342        assert_eq!(rate, 0.0);
343    }
344
345    #[test]
346    fn test_get_events_in_range_partial() {
347        let event_store = Arc::new(EventStore::new());
348        let event1 = MemoryEvent::allocate(0x1000, 1024, 100);
349        let event2 = MemoryEvent::allocate(0x2000, 2048, 200);
350        let event3 = MemoryEvent::allocate(0x3000, 4096, 300);
351        event_store.record(event1);
352        event_store.record(event2);
353        event_store.record(event3);
354
355        let query = TimelineQuery::new(event_store);
356        // Use a wide range to ensure we capture events
357        let events = query.get_events_in_range(0, u64::MAX);
358        assert!(!events.is_empty());
359    }
360
361    #[test]
362    fn test_memory_usage_over_time_with_deallocations() {
363        let event_store = Arc::new(EventStore::new());
364        let alloc = MemoryEvent::allocate(0x1000, 1024, 100);
365        let dealloc = MemoryEvent::deallocate(0x1000, 1024, 500);
366        event_store.record(alloc);
367        event_store.record(dealloc);
368
369        let query = TimelineQuery::new(event_store);
370        let usage = query.get_memory_usage_over_time(0, 1000, 100);
371
372        // Verify we get usage data
373        assert!(!usage.is_empty());
374    }
375
376    #[test]
377    fn test_memory_usage_over_time_with_reallocations() {
378        let event_store = Arc::new(EventStore::new());
379        let realloc = MemoryEvent::reallocate(0x1000, 1024, 2048, 200);
380        event_store.record(realloc);
381
382        let query = TimelineQuery::new(event_store);
383        let usage = query.get_memory_usage_over_time(0, 1000, 100);
384        assert!(!usage.is_empty());
385    }
386
387    #[test]
388    fn test_get_peak_memory_with_reallocations() {
389        let event_store = Arc::new(EventStore::new());
390        let alloc = MemoryEvent::allocate(0x1000, 1024, 100);
391        let realloc = MemoryEvent::reallocate(0x1000, 1024, 2048, 200);
392        event_store.record(alloc);
393        event_store.record(realloc);
394
395        let query = TimelineQuery::new(event_store);
396        let peak = query.get_peak_memory_in_range(0, u64::MAX);
397        // Peak should be 2048 (after reallocation)
398        assert_eq!(peak, 2048);
399    }
400
401    #[test]
402    fn test_get_events_in_range_boundary() {
403        let event_store = Arc::new(EventStore::new());
404        let event1 = MemoryEvent::allocate(0x1000, 1024, 100);
405        let event2 = MemoryEvent::allocate(0x2000, 2048, 200);
406        event_store.record(event1);
407        event_store.record(event2);
408
409        let query = TimelineQuery::new(event_store);
410
411        // Use a wide range to ensure we capture events
412        let events_all = query.get_events_in_range(0, u64::MAX);
413        assert!(!events_all.is_empty());
414    }
415
416    #[test]
417    fn test_get_allocations_empty_store() {
418        let event_store = Arc::new(EventStore::new());
419        let query = TimelineQuery::new(event_store);
420        let allocations = query.get_allocations_in_range(0, 1000);
421        assert!(allocations.is_empty());
422    }
423
424    #[test]
425    fn test_get_deallocations_empty_store() {
426        let event_store = Arc::new(EventStore::new());
427        let query = TimelineQuery::new(event_store);
428        let deallocations = query.get_deallocations_in_range(0, 1000);
429        assert!(deallocations.is_empty());
430    }
431}