dfir_rs/scheduled/
metrics.rs

1//! Runtime metrics for DFIR.
2
3use std::cell::Cell;
4use std::pin::Pin;
5use std::rc::Rc;
6use std::task::{Context, Poll};
7
8use dfir_lang::graph_ids::{GraphNodeId, GraphSubgraphId};
9use pin_project_lite::pin_project;
10use sealed::sealed;
11use slotmap::SecondaryMap;
12use web_time::{Duration, Instant};
13
14/// Typestate marker for whether a [`Context`](super::context::Context) tracks runtime metrics.
15///
16/// This trait is sealed and is implemented for exactly two types:
17///
18/// - [`DfirMetrics`]: metrics are tracked. The context holds a shared [`Rc<DfirMetrics>`] handle
19///   which is updated by the generated code and exposed via
20///   [`Dfir::metrics`](super::context::Dfir::metrics).
21/// - `()`: metrics are not tracked. The handle is the zero-sized `()`, so the context statically
22///   contains no metrics state at all, and the `metrics()` accessors do not exist.
23///
24/// Which one is used is determined at codegen time by
25/// [`AsCodeOptions::include_metrics_tracking`](dfir_lang::graph::AsCodeOptions::include_metrics_tracking),
26/// so there is no runtime check or default-allocated metrics when tracking is disabled.
27#[sealed]
28pub trait MetricsState: 'static {
29    /// The metrics handle stored inside the context: [`Rc<DfirMetrics>`] or `()`.
30    type Handle;
31}
32
33/// Metrics are tracked.
34#[sealed]
35impl MetricsState for DfirMetrics {
36    type Handle = Rc<DfirMetrics>;
37}
38
39/// Metrics are not tracked. Zero-sized, no runtime cost.
40#[sealed]
41impl MetricsState for () {
42    type Handle = ();
43}
44
45/// Metrics for a [`Dfir`](super::context::Dfir) graph instance.
46///
47/// Call [`Dfir::metrics`](super::context::Dfir::metrics) for reference-counted continually-updated metrics,
48/// or call [`Dfir::metrics_intervals`](super::context::Dfir::metrics_intervals) to obtain a [`DfirMetricsIntervals`] handle, and use
49/// [`DfirMetricsIntervals::take_interval`] to retrieve metrics for successive intervals.
50#[derive(Default, Clone)]
51#[non_exhaustive]
52pub struct DfirMetrics {
53    /// Per-subgraph metrics.
54    pub subgraphs: SecondaryMap<GraphSubgraphId, SubgraphMetrics>,
55    /// Per-handoff metrics.
56    pub handoffs: SecondaryMap<GraphNodeId, HandoffMetrics>,
57}
58
59impl DfirMetrics {
60    /// Subtracts `other` from self.
61    pub(super) fn diff(&mut self, other: &Self) {
62        for (sg_id, prev_sg_metrics) in other.subgraphs.iter() {
63            if let Some(curr_sg_metrics) = self.subgraphs.get_mut(sg_id) {
64                curr_sg_metrics.diff(prev_sg_metrics);
65            }
66        }
67        for (handoff_id, prev_handoff_metrics) in other.handoffs.iter() {
68            if let Some(curr_handoff_metrics) = self.handoffs.get_mut(handoff_id) {
69                curr_handoff_metrics.diff(prev_handoff_metrics);
70            }
71        }
72    }
73}
74
75/// A handle into a DFIR instance's metrics, where each call to [`Self::take_interval`] ends the current interval and
76/// returns its metrics. Obtained via [`Dfir::metrics_intervals`](super::context::Dfir::metrics_intervals).
77///
78/// The first call to `take_interval` returns metrics since this DFIR instance was created. Each subsequent call to
79/// `take_interval` returns metrics since the previous call.
80///
81/// Cloning the handle "forks" it from the original, as afterwards each interval may return different metrics
82/// depending on when exactly `take_interval` is called.
83#[derive(Clone)]
84pub struct DfirMetricsIntervals {
85    /// `curr` is continually updating (via shared ownership).
86    pub(super) curr: Rc<DfirMetrics>,
87    /// `prev` is an unchanging snapshot in time. `None` for "since creation".
88    pub(super) prev: Option<DfirMetrics>,
89}
90
91impl DfirMetricsIntervals {
92    /// Ends the current interval and returns the accumulated metrics across the interval.
93    ///
94    /// The first call to `take_interval` returns metrics since this DFIR instance was created. Each subsequent call to
95    /// `take_interval` returns metrics since the previous call.
96    pub fn take_interval(&mut self) -> DfirMetrics {
97        let mut curr = self.curr.as_ref().clone();
98        if let Some(prev) = self.prev.replace(curr.clone()) {
99            curr.diff(&prev);
100        }
101        curr
102    }
103
104    /// Returns a reference-counted handle to the original continually-updated runtime metrics for this DFIR instance.
105    ///
106    /// See [`Dfir::metrics`](super::context::Dfir::metrics).
107    pub fn all_metrics(&self) -> Rc<DfirMetrics> {
108        Rc::clone(&self.curr)
109    }
110}
111
112/// Declarative macro to generate metrics structs with Cell-based fields and getter methods.
113macro_rules! define_metrics {
114    (
115        $(#[$struct_attr:meta])*
116        pub struct $struct_name:ident {
117            $(
118                $( #[doc = $doc:literal] )*
119                #[diff($diff:ident)]
120                $( #[$field_attr:meta] )*
121                $field_name:ident: Cell<$field_type:ty>,
122            )*
123        }
124    ) => {
125        $(#[$struct_attr])*
126        #[derive(Default, Debug, Clone)]
127        #[non_exhaustive] // May add more metrics later.
128        pub struct $struct_name {
129            $(
130                #[doc(hidden)] // Public for codegen access; use the getter method instead.
131                $(#[$field_attr])*
132                pub $field_name: Cell<$field_type>,
133            )*
134        }
135
136        impl $struct_name {
137            $(
138                $( #[doc = $doc] )*
139                pub fn $field_name(&self) -> $field_type {
140                    self.$field_name.get()
141                }
142            )*
143
144            fn diff(&mut self, other: &Self) {
145                $(
146                    define_metrics_diff_field!($diff, $field_name, self, other);
147                )*
148            }
149        }
150    };
151}
152
153macro_rules! define_metrics_diff_field {
154    (total, $field:ident, $slf:ident, $other:ident) => {
155        debug_assert!($other.$field.get() <= $slf.$field.get());
156        $slf.$field.update(|x| x - $other.$field.get());
157    };
158    (curr, $field:ident, $slf:ident, $other:ident) => {};
159}
160
161define_metrics! {
162    /// Per-handoff metrics.
163    pub struct HandoffMetrics {
164        /// Number of items currently in the handoff.
165        #[diff(curr)]
166        curr_items_count: Cell<usize>,
167
168        /// Total number of items read out of the handoff.
169        #[diff(total)]
170        total_items_count: Cell<usize>,
171    }
172}
173
174define_metrics! {
175    /// Per-subgraph metrics.
176    pub struct SubgraphMetrics {
177        /// Number of times the subgraph has run.
178        #[diff(total)]
179        total_run_count: Cell<usize>,
180
181        /// Time elapsed during polling (when the subgraph is actively doing work).
182        #[diff(total)]
183        total_poll_duration: Cell<Duration>,
184
185        /// Number of times the subgraph has been polled.
186        #[diff(total)]
187        total_poll_count: Cell<usize>,
188
189        /// Time elapsed during idle (when the subgraph has yielded and is waiting for async events).
190        #[diff(total)]
191        total_idle_duration: Cell<Duration>,
192
193        /// Number of times the subgraph has been idle.
194        #[diff(total)]
195        total_idle_count: Cell<usize>,
196    }
197}
198
199pin_project! {
200    /// Helper struct which instruments a future to track polling times.
201    #[doc(hidden)]
202    pub struct InstrumentSubgraph<'a, Fut> {
203        #[pin]
204        future: Fut,
205        idle_start: Option<Instant>,
206        metrics: &'a SubgraphMetrics,
207    }
208}
209
210impl<'a, Fut> InstrumentSubgraph<'a, Fut> {
211    /// Wrap a future to track per-subgraph poll and idle durations.
212    pub fn new(future: Fut, metrics: &'a SubgraphMetrics) -> Self {
213        Self {
214            future,
215            idle_start: None,
216            metrics,
217        }
218    }
219}
220
221impl<'a, Fut> Future for InstrumentSubgraph<'a, Fut>
222where
223    Fut: Future,
224{
225    type Output = Fut::Output;
226    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
227        let this = self.project();
228
229        // End idle duration.
230        if let Some(idle_start) = this.idle_start {
231            this.metrics
232                .total_idle_duration
233                .update(|x| x + idle_start.elapsed());
234            this.metrics.total_idle_count.update(|x| x + 1);
235        }
236
237        // Begin poll duration.
238        let poll_start = Instant::now();
239        let out = this.future.poll(cx);
240
241        // End poll duration.
242        this.metrics
243            .total_poll_duration
244            .update(|x| x + poll_start.elapsed());
245        this.metrics.total_poll_count.update(|x| x + 1);
246
247        // Begin idle duration.
248        this.idle_start.replace(Instant::now());
249
250        out
251    }
252}
253
254#[cfg(test)]
255mod test {
256    use dfir_lang::graph_ids::{GraphNodeId, GraphSubgraphId};
257    use slotmap::SlotMap;
258
259    use super::*;
260
261    #[test]
262    fn test_dfir_metrics_intervals() {
263        // Create slotmaps to generate valid keys.
264        let mut sg_map: SlotMap<GraphSubgraphId, ()> = SlotMap::with_key();
265        let mut node_map: SlotMap<GraphNodeId, ()> = SlotMap::with_key();
266        let sg_id = sg_map.insert(());
267        let handoff_id = node_map.insert(());
268
269        let mut metrics = DfirMetrics::default();
270        metrics.subgraphs.insert(
271            sg_id,
272            SubgraphMetrics {
273                total_run_count: Cell::new(5),
274                total_poll_count: Cell::new(10),
275                total_idle_count: Cell::new(2),
276                total_poll_duration: Cell::new(Duration::from_millis(500)),
277                total_idle_duration: Cell::new(Duration::from_millis(200)),
278            },
279        );
280        metrics.handoffs.insert(
281            handoff_id,
282            HandoffMetrics {
283                curr_items_count: Cell::new(3),
284                total_items_count: Cell::new(100),
285            },
286        );
287        let metrics = Rc::new(metrics);
288
289        let mut intervals = DfirMetricsIntervals {
290            curr: Rc::clone(&metrics),
291            prev: None,
292        };
293
294        // First iteration - captures initial state
295        let first = intervals.take_interval();
296        let sg_metrics = &first.subgraphs[sg_id];
297        assert_eq!(sg_metrics.total_run_count(), 5);
298        let hoff_metrics = &first.handoffs[handoff_id];
299        assert_eq!(hoff_metrics.total_items_count(), 100);
300        assert_eq!(hoff_metrics.curr_items_count(), 3);
301
302        // Simulate more work being done.
303        let sg_metrics = &metrics.subgraphs[sg_id];
304        sg_metrics.total_run_count.set(12);
305        sg_metrics.total_poll_count.set(25);
306        sg_metrics.total_idle_count.set(7);
307        sg_metrics
308            .total_poll_duration
309            .set(Duration::from_millis(1200));
310        sg_metrics
311            .total_idle_duration
312            .set(Duration::from_millis(600));
313        let hoff_metrics = &metrics.handoffs[handoff_id];
314        hoff_metrics.total_items_count.set(250);
315        hoff_metrics.curr_items_count.set(10);
316
317        // Second iteration - should return the diff
318        let second = intervals.take_interval();
319        let sg_metrics = &second.subgraphs[sg_id];
320        assert_eq!(sg_metrics.total_run_count(), 7); // 12 - 5
321        assert_eq!(sg_metrics.total_poll_count(), 15); // 25 - 10
322        assert_eq!(sg_metrics.total_idle_count(), 5); // 7 - 2
323        //
324        let hoff_metrics = &second.handoffs[handoff_id];
325        // total_items_count should be diffed
326        assert_eq!(hoff_metrics.total_items_count(), 150); // 250 - 100
327        // curr_items_count should NOT be diffed (it's a current value, not cumulative)
328        assert_eq!(hoff_metrics.curr_items_count(), 10);
329    }
330}
dfir_rs/scheduled/metrics.rs

dfir_rs/scheduled/
metrics.rs
