micromegas_tracing/event/
in_memory_sink.rs

1use super::{EventSink, StreamDesc, TracingBlock};
2use crate::{
3    images::{ImageBlock, ImageStream},
4    logs::{LogBlock, LogMetadata, LogStream},
5    metrics::{MetricsBlock, MetricsStream},
6    prelude::*,
7    property_set::Property,
8    spans::{ThreadBlock, ThreadStream},
9};
10use std::{
11    fmt,
12    sync::{Arc, Mutex},
13};
14
15pub struct MemSinkState {
16    pub process_info: Option<Arc<ProcessInfo>>,
17    pub image_stream_desc: Option<Arc<StreamDesc>>,
18    pub image_blocks: Vec<Arc<ImageBlock>>,
19    pub log_stream_desc: Option<Arc<StreamDesc>>,
20    pub metrics_stream_desc: Option<Arc<StreamDesc>>,
21    pub thread_stream_descs: Vec<Arc<StreamDesc>>,
22    pub thread_blocks: Vec<Arc<ThreadBlock>>,
23    pub log_blocks: Vec<Arc<LogBlock>>,
24    pub metrics_blocks: Vec<Arc<MetricsBlock>>,
25}
26
27/// for tests where we want to inspect the collected data
28pub struct InMemorySink {
29    pub state: Mutex<MemSinkState>,
30}
31
32impl InMemorySink {
33    pub fn new() -> Self {
34        let state = MemSinkState {
35            process_info: None,
36            image_stream_desc: None,
37            image_blocks: vec![],
38            log_stream_desc: None,
39            metrics_stream_desc: None,
40            thread_stream_descs: vec![],
41            thread_blocks: vec![],
42            log_blocks: vec![],
43            metrics_blocks: vec![],
44        };
45        Self {
46            state: Mutex::new(state),
47        }
48    }
49}
50
51impl Default for InMemorySink {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57impl EventSink for InMemorySink {
58    fn on_startup(&self, process_info: Arc<ProcessInfo>) {
59        self.state.lock().unwrap().process_info = Some(process_info);
60    }
61
62    fn on_shutdown(&self) {}
63
64    fn on_log_enabled(&self, _metadata: &LogMetadata) -> bool {
65        true // Enable all log events for testing
66    }
67
68    fn on_log(
69        &self,
70        _desc: &LogMetadata,
71        _properties: &[Property],
72        _time: i64,
73        _args: fmt::Arguments<'_>,
74    ) {
75        // For testing, we primarily collect events through blocks
76        // Individual log events are handled via on_process_log_block
77    }
78
79    fn on_init_log_stream(&self, log_stream: &LogStream) {
80        self.state.lock().unwrap().log_stream_desc = Some(log_stream.desc());
81    }
82
83    fn on_process_log_block(&self, log_block: Arc<LogBlock>) {
84        self.state.lock().unwrap().log_blocks.push(log_block);
85    }
86
87    fn on_init_metrics_stream(&self, metrics_stream: &MetricsStream) {
88        self.state.lock().unwrap().metrics_stream_desc = Some(metrics_stream.desc());
89    }
90
91    fn on_process_metrics_block(&self, metrics_block: Arc<MetricsBlock>) {
92        self.state
93            .lock()
94            .unwrap()
95            .metrics_blocks
96            .push(metrics_block);
97    }
98
99    fn on_init_image_stream(&self, image_stream: &ImageStream) {
100        self.state.lock().unwrap().image_stream_desc = Some(image_stream.desc());
101    }
102
103    fn on_process_image_block(&self, image_block: Arc<ImageBlock>) {
104        self.state.lock().unwrap().image_blocks.push(image_block);
105    }
106
107    fn on_init_thread_stream(&self, thread_stream: &ThreadStream) {
108        self.state
109            .lock()
110            .unwrap()
111            .thread_stream_descs
112            .push(thread_stream.desc());
113    }
114
115    fn on_process_thread_block(&self, thread_block: Arc<ThreadBlock>) {
116        self.state.lock().unwrap().thread_blocks.push(thread_block);
117    }
118
119    fn is_busy(&self) -> bool {
120        false // For testing, never report as busy
121    }
122}
123
124impl InMemorySink {
125    /// Get the total number of image blocks collected
126    pub fn image_block_count(&self) -> usize {
127        self.state.lock().unwrap().image_blocks.len()
128    }
129
130    /// Get the total number of thread blocks collected
131    pub fn thread_block_count(&self) -> usize {
132        self.state.lock().unwrap().thread_blocks.len()
133    }
134
135    /// Get the total number of log blocks collected
136    pub fn log_block_count(&self) -> usize {
137        self.state.lock().unwrap().log_blocks.len()
138    }
139
140    /// Get the total number of metrics blocks collected
141    pub fn metrics_block_count(&self) -> usize {
142        self.state.lock().unwrap().metrics_blocks.len()
143    }
144
145    /// Get the total number of events across all thread blocks
146    pub fn total_thread_events(&self) -> usize {
147        self.state
148            .lock()
149            .unwrap()
150            .thread_blocks
151            .iter()
152            .map(|block| block.nb_objects())
153            .sum()
154    }
155
156    /// Get the total number of events across all log blocks
157    pub fn total_log_events(&self) -> usize {
158        self.state
159            .lock()
160            .unwrap()
161            .log_blocks
162            .iter()
163            .map(|block| block.nb_objects())
164            .sum()
165    }
166
167    /// Get the total number of events across all metrics blocks
168    pub fn total_metrics_events(&self) -> usize {
169        self.state
170            .lock()
171            .unwrap()
172            .metrics_blocks
173            .iter()
174            .map(|block| block.nb_objects())
175            .sum()
176    }
177}
micromegas_tracing/event/in_memory_sink.rs

micromegas_tracing/event/
in_memory_sink.rs
