Skip to main content

datum/graph/
metrics.rs

1//! Per-node fused-graph metrics and the disabled-path probe seam.
2
3use super::FusedNodeAttributes;
4use crate::MetricsLevel;
5use std::{
6    sync::atomic::{AtomicBool, Ordering},
7    time::{Duration, Instant},
8};
9
10/// Immutable metrics for one fused graph node.
11///
12/// Scalar stages populate `elements_*`. Arrow batch stages populate `batches_*` and `rows_*`.
13/// Typed fused kernels measure a whole execution chunk with a monotonic clock and publish an equal
14/// share to the participating nodes; erased nodes measure their individual transitions. Stall time
15/// is measured exactly around bounded-buffer/demand waits. Counts are always exact.
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct FusedNodeMetrics {
18    pub stage_index: usize,
19    pub stage_name: String,
20    pub effective_name: String,
21    pub metrics_level: MetricsLevel,
22    pub elements_in: u64,
23    pub elements_out: u64,
24    pub batches_in: u64,
25    pub batches_out: u64,
26    pub rows_in: u64,
27    pub rows_out: u64,
28    pub processing_time: Duration,
29    pub backpressure_stall_time: Duration,
30    pub queue_depth: Option<usize>,
31}
32
33#[derive(Clone, Copy, Debug, Default)]
34pub(super) struct MetricUnits {
35    pub(super) elements: u64,
36    pub(super) batches: u64,
37    pub(super) rows: u64,
38}
39
40impl MetricUnits {
41    pub(super) const fn scalar(elements: u64) -> Self {
42        Self {
43            elements,
44            batches: 0,
45            rows: 0,
46        }
47    }
48
49    #[allow(dead_code)]
50    pub(super) const fn batches(batches: u64, rows: u64) -> Self {
51        Self {
52            elements: 0,
53            batches,
54            rows,
55        }
56    }
57}
58
59#[derive(Debug)]
60struct NodeState {
61    report: FusedNodeMetrics,
62    stall_started: Option<Instant>,
63}
64
65/// Fresh per-materialization state. Disabled runs are stack-only and own no node allocation.
66#[derive(Debug)]
67pub(super) struct FusedMetricsRun {
68    enabled: AtomicBool,
69    nodes: Vec<NodeState>,
70}
71
72#[derive(Debug)]
73pub(super) enum NodeMeasurement {
74    Disabled,
75    Counts { stage_index: usize },
76    Timing { stage_index: usize, start: Instant },
77}
78
79#[derive(Debug)]
80pub(super) enum StallMeasurement {
81    Disabled,
82    Timing { stage_index: usize, start: Instant },
83}
84
85impl FusedMetricsRun {
86    pub(super) fn disabled() -> Self {
87        Self {
88            enabled: AtomicBool::new(false),
89            nodes: Vec::new(),
90        }
91    }
92
93    pub(super) fn new(attributes: &[FusedNodeAttributes]) -> Self {
94        let enabled = attributes.iter().any(|node| {
95            node.attributes
96                .metrics_level_hint()
97                .unwrap_or(MetricsLevel::Off)
98                != MetricsLevel::Off
99        });
100        if !enabled {
101            return Self::disabled();
102        }
103
104        let nodes = attributes
105            .iter()
106            .map(|node| NodeState {
107                report: FusedNodeMetrics {
108                    stage_index: node.stage_index,
109                    stage_name: node.stage_name.clone(),
110                    effective_name: node.effective_name.clone(),
111                    metrics_level: node
112                        .attributes
113                        .metrics_level_hint()
114                        .unwrap_or(MetricsLevel::Off),
115                    elements_in: 0,
116                    elements_out: 0,
117                    batches_in: 0,
118                    batches_out: 0,
119                    rows_in: 0,
120                    rows_out: 0,
121                    processing_time: Duration::ZERO,
122                    backpressure_stall_time: Duration::ZERO,
123                    queue_depth: None,
124                },
125                stall_started: None,
126            })
127            .collect();
128        Self {
129            enabled: AtomicBool::new(true),
130            nodes,
131        }
132    }
133
134    /// One relaxed load used to select the untouched hot loop or its instrumented twin.
135    #[inline(always)]
136    pub(super) fn enabled(&self) -> bool {
137        metrics_gate_enabled(&self.enabled)
138    }
139
140    pub(super) fn timing_enabled(&self) -> bool {
141        self.nodes.iter().any(|node| {
142            matches!(
143                node.report.metrics_level,
144                MetricsLevel::Timing | MetricsLevel::Stalls
145            )
146        })
147    }
148
149    pub(super) fn complete_exact_counts(
150        &mut self,
151        stage_index: usize,
152        input: MetricUnits,
153        output: MetricUnits,
154    ) {
155        let Some(node) = self.nodes.get_mut(stage_index) else {
156            return;
157        };
158        if node.report.metrics_level == MetricsLevel::Off {
159            return;
160        }
161        node.report.elements_in = input.elements;
162        node.report.elements_out = output.elements;
163        node.report.batches_in = input.batches;
164        node.report.batches_out = output.batches;
165        node.report.rows_in = input.rows;
166        node.report.rows_out = output.rows;
167    }
168
169    pub(super) fn complete_linear_metrics(
170        &mut self,
171        elements: u64,
172        processing_time: Option<Duration>,
173    ) {
174        let timed_nodes = self
175            .nodes
176            .iter()
177            .filter(|node| {
178                matches!(
179                    node.report.metrics_level,
180                    MetricsLevel::Timing | MetricsLevel::Stalls
181                )
182            })
183            .count();
184        let processing_share = processing_time.and_then(|elapsed| {
185            (timed_nodes != 0).then(|| {
186                let nanos = elapsed.as_nanos() / timed_nodes as u128;
187                Duration::from_nanos(nanos.min(u128::from(u64::MAX)) as u64)
188            })
189        });
190
191        for node in &mut self.nodes {
192            if node.report.metrics_level == MetricsLevel::Off {
193                continue;
194            }
195            node.report.elements_in = elements;
196            node.report.elements_out = elements;
197            if let Some(share) = processing_share
198                && matches!(
199                    node.report.metrics_level,
200                    MetricsLevel::Timing | MetricsLevel::Stalls
201                )
202            {
203                node.report.processing_time = share;
204            }
205        }
206    }
207
208    pub(super) fn record_processing_time_share(
209        &mut self,
210        stage_indices: &[usize],
211        elapsed: Duration,
212    ) {
213        let timed_nodes = stage_indices
214            .iter()
215            .filter(|&&stage_index| {
216                self.nodes.get(stage_index).is_some_and(|node| {
217                    matches!(
218                        node.report.metrics_level,
219                        MetricsLevel::Timing | MetricsLevel::Stalls
220                    )
221                })
222            })
223            .count();
224        if timed_nodes == 0 {
225            return;
226        }
227        let share = elapsed.as_nanos() / timed_nodes as u128;
228        let share = Duration::from_nanos(share.min(u128::from(u64::MAX)) as u64);
229        for &stage_index in stage_indices {
230            let Some(node) = self.nodes.get_mut(stage_index) else {
231                continue;
232            };
233            if matches!(
234                node.report.metrics_level,
235                MetricsLevel::Timing | MetricsLevel::Stalls
236            ) {
237                node.report.processing_time = node.report.processing_time.saturating_add(share);
238            }
239        }
240    }
241
242    /// Enter an enabled node site after the relaxed loop-entry gate has selected instrumentation.
243    #[inline(always)]
244    pub(super) fn begin_node(&mut self, stage_index: usize, input: MetricUnits) -> NodeMeasurement {
245        let Some(node) = self.nodes.get_mut(stage_index) else {
246            return NodeMeasurement::Disabled;
247        };
248        if node.report.metrics_level == MetricsLevel::Off {
249            return NodeMeasurement::Disabled;
250        }
251        node.report.elements_in = node.report.elements_in.saturating_add(input.elements);
252        node.report.batches_in = node.report.batches_in.saturating_add(input.batches);
253        node.report.rows_in = node.report.rows_in.saturating_add(input.rows);
254
255        match node.report.metrics_level {
256            MetricsLevel::Off => NodeMeasurement::Disabled,
257            MetricsLevel::Counts => NodeMeasurement::Counts { stage_index },
258            MetricsLevel::Timing | MetricsLevel::Stalls => NodeMeasurement::Timing {
259                stage_index,
260                start: Instant::now(),
261            },
262        }
263    }
264
265    #[inline(always)]
266    pub(super) fn end_node(&mut self, measurement: NodeMeasurement, output: MetricUnits) {
267        let (stage_index, elapsed) = match measurement {
268            NodeMeasurement::Disabled => return,
269            NodeMeasurement::Counts { stage_index } => (stage_index, None),
270            NodeMeasurement::Timing { stage_index, start } => (stage_index, Some(start.elapsed())),
271        };
272        let node = &mut self.nodes[stage_index].report;
273        node.elements_out = node.elements_out.saturating_add(output.elements);
274        node.batches_out = node.batches_out.saturating_add(output.batches);
275        node.rows_out = node.rows_out.saturating_add(output.rows);
276        if let Some(elapsed) = elapsed {
277            node.processing_time = node.processing_time.saturating_add(elapsed);
278        }
279    }
280
281    /// Observe a bounded stage transition. A no-output transition begins a stall; the next output
282    /// closes it. This covers demand/peer waits in the erased bounded-buffer state machines.
283    pub(super) fn observe_backpressure(&mut self, stage_index: usize, output_count: usize) {
284        let Some(node) = self.nodes.get_mut(stage_index) else {
285            return;
286        };
287        if node.report.metrics_level != MetricsLevel::Stalls {
288            return;
289        }
290        if output_count == 0 {
291            if node.stall_started.is_none() {
292                node.stall_started = Some(Instant::now());
293            }
294            node.report.queue_depth = Some(node.report.queue_depth.unwrap_or(0).max(1));
295        } else if let Some(started) = node.stall_started.take() {
296            node.report.backpressure_stall_time = node
297                .report
298                .backpressure_stall_time
299                .saturating_add(started.elapsed());
300        }
301    }
302
303    /// Enter a bounded-buffer/demand wait site.
304    #[inline(always)]
305    #[allow(dead_code)]
306    pub(super) fn begin_stall(&self, stage_index: usize) -> StallMeasurement {
307        if !self.enabled.load(Ordering::Relaxed) {
308            return StallMeasurement::Disabled;
309        }
310        if self
311            .nodes
312            .get(stage_index)
313            .is_none_or(|node| node.report.metrics_level != MetricsLevel::Stalls)
314        {
315            return StallMeasurement::Disabled;
316        }
317        StallMeasurement::Timing {
318            stage_index,
319            start: Instant::now(),
320        }
321    }
322
323    #[allow(dead_code)]
324    pub(super) fn end_stall(&mut self, measurement: StallMeasurement, queue_depth: Option<usize>) {
325        let StallMeasurement::Timing { stage_index, start } = measurement else {
326            return;
327        };
328        let node = &mut self.nodes[stage_index].report;
329        node.backpressure_stall_time = node.backpressure_stall_time.saturating_add(start.elapsed());
330        if let Some(depth) = queue_depth {
331            node.queue_depth = Some(node.queue_depth.unwrap_or(0).max(depth));
332        }
333    }
334
335    pub(super) fn snapshot(&self) -> Vec<FusedNodeMetrics> {
336        if !self.enabled.load(Ordering::Relaxed) {
337            return Vec::new();
338        }
339        self.nodes
340            .iter()
341            .filter(|node| node.report.metrics_level != MetricsLevel::Off)
342            .map(|node| node.report.clone())
343            .collect()
344    }
345}
346
347/// Stable, non-inlined shim used to inspect the disabled-path generated code.
348#[doc(hidden)]
349#[inline(never)]
350pub fn metrics_disabled_path_probe(flag: &AtomicBool, value: u64) -> u64 {
351    if !metrics_gate_enabled(flag) {
352        return value.wrapping_mul(3).wrapping_add(1);
353    }
354    metrics_enabled_path_probe(value)
355}
356
357/// Same-signature no-site control for path-sensitive generated-code comparison.
358#[doc(hidden)]
359#[inline(never)]
360pub fn metrics_no_site_probe(_flag: &AtomicBool, value: u64) -> u64 {
361    value.wrapping_mul(3).wrapping_add(1)
362}
363
364#[inline(always)]
365fn metrics_gate_enabled(flag: &AtomicBool) -> bool {
366    if flag.load(Ordering::Relaxed) {
367        return metrics_gate_enabled_cold();
368    }
369    false
370}
371
372#[cold]
373#[inline(never)]
374fn metrics_gate_enabled_cold() -> bool {
375    true
376}
377
378#[cold]
379#[inline(never)]
380fn metrics_enabled_path_probe(value: u64) -> u64 {
381    value.wrapping_mul(3).wrapping_add(2)
382}