1use crate::event_store::event::{MemoryEvent, MemoryEventType};
7use crate::event_store::EventStore;
8use std::sync::Arc;
9
10pub struct TimelineQuery {
15 event_store: Arc<EventStore>,
17}
18
19impl TimelineQuery {
20 pub fn new(event_store: Arc<EventStore>) -> Self {
22 Self { event_store }
23 }
24
25 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 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 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 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 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 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 all_events.sort_by_key(|(ts, _)| *ts);
108
109 let mut current = start;
111 let mut event_idx = 0;
112 let mut running_memory: usize = 0;
113
114 while current < end {
115 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 pub fn get_peak_memory_in_range(&self, start: u64, end: u64) -> usize {
158 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 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 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 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 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 let rate = query.get_event_rate(0, 1000);
330 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 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 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 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 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}