memscope_rs/query/
engine.rs

1//! Query Engine - Unified query interface
2//!
3//! This module provides the QueryEngine which is responsible for
4//! querying snapshot data.
5
6use crate::query::types::{
7    AllocationQueryResult, QueryResult, SummaryQueryResult, ThreadQueryResult,
8};
9use crate::snapshot::{MemorySnapshot, SharedSnapshotEngine};
10use std::sync::Arc;
11
12/// Query Engine - Unified query interface for memory data
13///
14/// The QueryEngine provides a unified interface for querying memory
15/// snapshot data, supporting various query types and filtering options.
16///
17/// Key properties:
18/// - Unified: Single interface for all query types
19/// - Efficient: Optimized for fast query execution
20/// - Flexible: Supports filtering and sorting
21pub struct QueryEngine {
22    /// Reference to the snapshot engine
23    snapshot_engine: SharedSnapshotEngine,
24}
25
26/// Shared reference to QueryEngine
27pub type SharedQueryEngine = Arc<QueryEngine>;
28
29impl QueryEngine {
30    /// Create a new QueryEngine
31    pub fn new(snapshot_engine: SharedSnapshotEngine) -> Self {
32        Self { snapshot_engine }
33    }
34
35    /// Get the current snapshot
36    fn get_snapshot(&self) -> MemorySnapshot {
37        self.snapshot_engine.build_snapshot()
38    }
39
40    /// Query for top allocations by size
41    ///
42    /// # Arguments
43    /// * `limit` - Maximum number of allocations to return
44    pub fn top_allocations(&self, limit: usize) -> QueryResult {
45        let snapshot = self.get_snapshot();
46        let mut allocations: Vec<_> = snapshot.active_allocations.values().cloned().collect();
47
48        // Sort by size descending
49        allocations.sort_by(|a, b| b.size.cmp(&a.size));
50
51        // Limit results
52        allocations.truncate(limit);
53
54        let total_bytes = allocations.iter().map(|a| a.size).sum();
55
56        QueryResult::Allocations(AllocationQueryResult {
57            count: allocations.len(),
58            total_bytes,
59            allocations,
60        })
61    }
62
63    /// Query for allocations from a specific thread
64    ///
65    /// # Arguments
66    /// * `thread_id` - The thread ID to filter by
67    pub fn allocations_by_thread(&self, thread_id: u64) -> QueryResult {
68        let snapshot = self.get_snapshot();
69        let allocations: Vec<_> = snapshot
70            .active_allocations
71            .values()
72            .filter(|a| a.thread_id == thread_id)
73            .cloned()
74            .collect();
75
76        let total_bytes = allocations.iter().map(|a| a.size).sum();
77
78        QueryResult::Allocations(AllocationQueryResult {
79            count: allocations.len(),
80            total_bytes,
81            allocations,
82        })
83    }
84
85    /// Query for thread statistics
86    pub fn thread_stats(&self) -> QueryResult {
87        let snapshot = self.get_snapshot();
88        let threads: Vec<_> = snapshot.thread_stats.values().cloned().collect();
89        let total_bytes = threads.iter().map(|t| t.current_memory).sum();
90
91        QueryResult::Threads(ThreadQueryResult {
92            count: threads.len(),
93            total_bytes,
94            threads,
95        })
96    }
97
98    /// Query for a summary of memory usage
99    pub fn summary(&self) -> QueryResult {
100        let snapshot = self.get_snapshot();
101
102        QueryResult::Summary(SummaryQueryResult {
103            total_allocations: snapshot.stats.total_allocations,
104            total_deallocations: snapshot.stats.total_deallocations,
105            active_allocations: snapshot.stats.active_allocations,
106            total_allocated: snapshot.stats.total_allocated,
107            total_deallocated: snapshot.stats.total_deallocated,
108            current_memory: snapshot.stats.current_memory,
109            peak_memory: snapshot.stats.peak_memory,
110            thread_count: snapshot.thread_stats.len(),
111        })
112    }
113
114    /// Query for allocations with a specific variable name
115    ///
116    /// # Arguments
117    /// * `var_name` - The variable name to filter by
118    pub fn allocations_by_variable(&self, var_name: &str) -> QueryResult {
119        let snapshot = self.get_snapshot();
120        let allocations: Vec<_> = snapshot
121            .active_allocations
122            .values()
123            .filter(|a| a.var_name.as_ref().map(|n| n == var_name).unwrap_or(false))
124            .cloned()
125            .collect();
126
127        let total_bytes = allocations.iter().map(|a| a.size).sum();
128
129        QueryResult::Allocations(AllocationQueryResult {
130            count: allocations.len(),
131            total_bytes,
132            allocations,
133        })
134    }
135
136    /// Query for allocations larger than a certain size
137    ///
138    /// # Arguments
139    /// * `min_size` - Minimum allocation size in bytes
140    pub fn allocations_larger_than(&self, min_size: usize) -> QueryResult {
141        let snapshot = self.get_snapshot();
142        let allocations: Vec<_> = snapshot
143            .active_allocations
144            .values()
145            .filter(|a| a.size > min_size)
146            .cloned()
147            .collect();
148
149        let total_bytes = allocations.iter().map(|a| a.size).sum();
150
151        QueryResult::Allocations(AllocationQueryResult {
152            count: allocations.len(),
153            total_bytes,
154            allocations,
155        })
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162    use crate::event_store::EventStore;
163    use crate::snapshot::SnapshotEngine;
164
165    /// Objective: Verify that QueryEngine creates correctly with empty data.
166    /// Invariants: Summary should show zero allocations for empty event store.
167    #[test]
168    fn test_query_engine_creation() {
169        let event_store = Arc::new(EventStore::new());
170        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
171        let query_engine = QueryEngine::new(snapshot_engine);
172
173        let result = query_engine.summary();
174        match result {
175            QueryResult::Summary(summary) => {
176                assert_eq!(
177                    summary.total_allocations, 0,
178                    "Empty event store should have zero allocations"
179                );
180            }
181            _ => panic!("Expected summary result"),
182        }
183    }
184
185    /// Objective: Verify that top_allocations returns largest allocations sorted by size.
186    /// Invariants: Results must be sorted descending by size and limited to requested count.
187    #[test]
188    fn test_top_allocations() {
189        let event_store = Arc::new(EventStore::new());
190        event_store.record(crate::event_store::MemoryEvent::allocate(0x1000, 1024, 1));
191        event_store.record(crate::event_store::MemoryEvent::allocate(0x2000, 2048, 1));
192        event_store.record(crate::event_store::MemoryEvent::allocate(0x3000, 512, 1));
193
194        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
195        let query_engine = QueryEngine::new(snapshot_engine);
196
197        let result = query_engine.top_allocations(2);
198        match result {
199            QueryResult::Allocations(allocations) => {
200                assert_eq!(allocations.count, 2, "Should return exactly 2 allocations");
201                assert_eq!(
202                    allocations.allocations[0].size, 2048,
203                    "First should be largest"
204                );
205                assert_eq!(
206                    allocations.allocations[1].size, 1024,
207                    "Second should be second largest"
208                );
209            }
210            _ => panic!("Expected allocations result"),
211        }
212    }
213
214    /// Objective: Verify that allocations_by_thread filters correctly by thread ID.
215    /// Invariants: Only allocations from specified thread should be returned.
216    #[test]
217    fn test_allocations_by_thread() {
218        let event_store = Arc::new(EventStore::new());
219        event_store.record(crate::event_store::MemoryEvent::allocate(0x1000, 1024, 1));
220        event_store.record(crate::event_store::MemoryEvent::allocate(0x2000, 2048, 2));
221
222        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
223        let query_engine = QueryEngine::new(snapshot_engine);
224
225        let result = query_engine.allocations_by_thread(1);
226        match result {
227            QueryResult::Allocations(allocations) => {
228                assert_eq!(
229                    allocations.count, 1,
230                    "Should have one allocation for thread 1"
231                );
232                assert_eq!(allocations.total_bytes, 1024, "Total bytes should be 1024");
233            }
234            _ => panic!("Expected allocations result"),
235        }
236    }
237
238    /// Objective: Verify that thread_stats returns correct thread statistics.
239    /// Invariants: Thread count should match number of unique threads with allocations.
240    #[test]
241    fn test_thread_stats() {
242        let event_store = Arc::new(EventStore::new());
243        event_store.record(crate::event_store::MemoryEvent::allocate(0x1000, 1024, 1));
244        event_store.record(crate::event_store::MemoryEvent::allocate(0x2000, 2048, 2));
245
246        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
247        let query_engine = QueryEngine::new(snapshot_engine);
248
249        let result = query_engine.thread_stats();
250        match result {
251            QueryResult::Threads(threads) => {
252                assert_eq!(threads.count, 2, "Should have 2 threads with allocations");
253            }
254            _ => panic!("Expected threads result"),
255        }
256    }
257
258    /// Objective: Verify that summary correctly tracks allocations and deallocations.
259    /// Invariants: Summary must accurately reflect allocation/deallocation counts.
260    #[test]
261    fn test_summary() {
262        let event_store = Arc::new(EventStore::new());
263        event_store.record(crate::event_store::MemoryEvent::allocate(0x1000, 1024, 1));
264        event_store.record(crate::event_store::MemoryEvent::deallocate(0x1000, 1024, 1));
265
266        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
267        let query_engine = QueryEngine::new(snapshot_engine);
268
269        let result = query_engine.summary();
270        match result {
271            QueryResult::Summary(summary) => {
272                assert_eq!(summary.total_allocations, 1, "Should have one allocation");
273                assert_eq!(
274                    summary.total_deallocations, 1,
275                    "Should have one deallocation"
276                );
277                assert_eq!(
278                    summary.active_allocations, 0,
279                    "Should have no active allocations"
280                );
281            }
282            _ => panic!("Expected summary result"),
283        }
284    }
285
286    /// Objective: Verify that allocations_by_variable filters correctly by variable name.
287    /// Invariants: Only allocations with matching var_name should be returned.
288    #[test]
289    fn test_allocations_by_variable() {
290        let event_store = Arc::new(EventStore::new());
291        let mut event1 = crate::event_store::MemoryEvent::allocate(0x1000, 1024, 1);
292        event1.var_name = Some("test_var".to_string());
293        event_store.record(event1);
294
295        let mut event2 = crate::event_store::MemoryEvent::allocate(0x2000, 2048, 1);
296        event2.var_name = Some("other_var".to_string());
297        event_store.record(event2);
298
299        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
300        let query_engine = QueryEngine::new(snapshot_engine);
301
302        let result = query_engine.allocations_by_variable("test_var");
303        match result {
304            QueryResult::Allocations(allocations) => {
305                assert_eq!(
306                    allocations.count, 1,
307                    "Should have one allocation with test_var"
308                );
309                assert_eq!(allocations.total_bytes, 1024, "Total bytes should be 1024");
310            }
311            _ => panic!("Expected allocations result"),
312        }
313    }
314
315    /// Objective: Verify that allocations_larger_than filters by size correctly.
316    /// Invariants: Only allocations larger than min_size should be returned.
317    #[test]
318    fn test_allocations_larger_than() {
319        let event_store = Arc::new(EventStore::new());
320        event_store.record(crate::event_store::MemoryEvent::allocate(0x1000, 100, 1));
321        event_store.record(crate::event_store::MemoryEvent::allocate(0x2000, 500, 1));
322        event_store.record(crate::event_store::MemoryEvent::allocate(0x3000, 1000, 1));
323
324        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
325        let query_engine = QueryEngine::new(snapshot_engine);
326
327        let result = query_engine.allocations_larger_than(200);
328        match result {
329            QueryResult::Allocations(allocations) => {
330                assert_eq!(
331                    allocations.count, 2,
332                    "Should have 2 allocations larger than 200"
333                );
334                assert_eq!(allocations.total_bytes, 1500, "Total bytes should be 1500");
335            }
336            _ => panic!("Expected allocations result"),
337        }
338    }
339
340    /// Objective: Verify that top_allocations with zero limit returns empty result.
341    /// Invariants: Zero limit should return empty allocations list.
342    #[test]
343    fn test_top_allocations_zero_limit() {
344        let event_store = Arc::new(EventStore::new());
345        event_store.record(crate::event_store::MemoryEvent::allocate(0x1000, 1024, 1));
346
347        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
348        let query_engine = QueryEngine::new(snapshot_engine);
349
350        let result = query_engine.top_allocations(0);
351        match result {
352            QueryResult::Allocations(allocations) => {
353                assert_eq!(
354                    allocations.count, 0,
355                    "Zero limit should return empty result"
356                );
357            }
358            _ => panic!("Expected allocations result"),
359        }
360    }
361
362    /// Objective: Verify that allocations_by_thread returns empty for unknown thread.
363    /// Invariants: Unknown thread ID should return empty allocations.
364    #[test]
365    fn test_allocations_by_unknown_thread() {
366        let event_store = Arc::new(EventStore::new());
367        event_store.record(crate::event_store::MemoryEvent::allocate(0x1000, 1024, 1));
368
369        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
370        let query_engine = QueryEngine::new(snapshot_engine);
371
372        let result = query_engine.allocations_by_thread(999);
373        match result {
374            QueryResult::Allocations(allocations) => {
375                assert_eq!(
376                    allocations.count, 0,
377                    "Unknown thread should have no allocations"
378                );
379            }
380            _ => panic!("Expected allocations result"),
381        }
382    }
383
384    /// Objective: Verify that allocations_by_variable returns empty for unknown variable.
385    /// Invariants: Unknown variable name should return empty allocations.
386    #[test]
387    fn test_allocations_by_unknown_variable() {
388        let event_store = Arc::new(EventStore::new());
389        let mut event = crate::event_store::MemoryEvent::allocate(0x1000, 1024, 1);
390        event.var_name = Some("known_var".to_string());
391        event_store.record(event);
392
393        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
394        let query_engine = QueryEngine::new(snapshot_engine);
395
396        let result = query_engine.allocations_by_variable("unknown_var");
397        match result {
398            QueryResult::Allocations(allocations) => {
399                assert_eq!(
400                    allocations.count, 0,
401                    "Unknown variable should have no allocations"
402                );
403            }
404            _ => panic!("Expected allocations result"),
405        }
406    }
407
408    /// Objective: Verify that allocations_larger_than returns empty when no matches.
409    /// Invariants: Very large min_size should return empty allocations.
410    #[test]
411    fn test_allocations_larger_than_no_matches() {
412        let event_store = Arc::new(EventStore::new());
413        event_store.record(crate::event_store::MemoryEvent::allocate(0x1000, 100, 1));
414
415        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
416        let query_engine = QueryEngine::new(snapshot_engine);
417
418        let result = query_engine.allocations_larger_than(10000);
419        match result {
420            QueryResult::Allocations(allocations) => {
421                assert_eq!(
422                    allocations.count, 0,
423                    "Very large min_size should return empty"
424                );
425            }
426            _ => panic!("Expected allocations result"),
427        }
428    }
429
430    /// Objective: Verify that summary tracks peak memory correctly.
431    /// Invariants: Peak memory should reflect maximum memory usage.
432    #[test]
433    fn test_summary_peak_memory() {
434        let event_store = Arc::new(EventStore::new());
435        event_store.record(crate::event_store::MemoryEvent::allocate(0x1000, 1024, 1));
436        event_store.record(crate::event_store::MemoryEvent::allocate(0x2000, 2048, 1));
437        event_store.record(crate::event_store::MemoryEvent::deallocate(0x1000, 1024, 1));
438
439        let snapshot_engine = Arc::new(SnapshotEngine::new(event_store));
440        let query_engine = QueryEngine::new(snapshot_engine);
441
442        let result = query_engine.summary();
443        match result {
444            QueryResult::Summary(summary) => {
445                assert!(
446                    summary.peak_memory >= 2048,
447                    "Peak memory should be at least 2048"
448                );
449            }
450            _ => panic!("Expected summary result"),
451        }
452    }
453}
memscope_rs/query/engine.rs

memscope_rs/query/
engine.rs
