openvm_stark_sdk/
metrics_tracing.rs

1use std::{sync::Arc, time::Instant};
2
3use dashmap::DashMap;
4use tracing::{
5    field::{Field, Visit},
6    Id, Subscriber,
7};
8use tracing_subscriber::{registry::LookupSpan, Layer};
9
10/// A tracing layer that automatically emits metric gauges for all span durations.
11/// This replaces the need for manual metrics_span calls by leveraging the tracing infrastructure.
12#[derive(Clone, Default)]
13pub struct TimingMetricsLayer {
14    /// Store span timings indexed by span ID
15    span_timings: Arc<DashMap<Id, SpanTiming>>,
16}
17
18#[derive(Debug)]
19struct SpanTiming {
20    name: String,
21    start_time: Instant,
22    labels: Vec<(String, String)>,
23}
24
25/// A visitor to extract the return value from span events
26struct ReturnValueVisitor {
27    has_return: bool,
28}
29
30impl Visit for ReturnValueVisitor {
31    fn record_debug(&mut self, field: &Field, _value: &dyn std::fmt::Debug) {
32        if field.name() == "return" {
33            self.has_return = true;
34        }
35    }
36
37    fn record_i64(&mut self, _field: &Field, _value: i64) {}
38    fn record_u64(&mut self, _field: &Field, _value: u64) {}
39    fn record_bool(&mut self, _field: &Field, _value: bool) {}
40    fn record_str(&mut self, _field: &Field, _value: &str) {}
41}
42
43/// A visitor to extract all string fields from span attributes as metric labels
44#[derive(Default)]
45struct LabelVisitor {
46    labels: Vec<(String, String)>,
47}
48
49impl Visit for LabelVisitor {
50    fn record_debug(&mut self, _field: &Field, _value: &dyn std::fmt::Debug) {}
51    fn record_i64(&mut self, _field: &Field, _value: i64) {}
52    fn record_u64(&mut self, _field: &Field, _value: u64) {}
53    fn record_bool(&mut self, _field: &Field, _value: bool) {}
54    fn record_str(&mut self, field: &Field, value: &str) {
55        self.labels
56            .push((field.name().to_string(), value.to_string()));
57    }
58}
59
60impl TimingMetricsLayer {
61    /// Create a new TimingMetricsLayer
62    pub fn new() -> Self {
63        Self::default()
64    }
65
66    fn emit_metric(name: &str, duration_ms: f64, labels: &[(String, String)]) {
67        let metric_name = format!("{}_time_ms", name);
68        let labels: Vec<metrics::Label> = labels
69            .iter()
70            .map(|(k, v)| metrics::Label::new(k.clone(), v.clone()))
71            .collect();
72        metrics::gauge!(metric_name, labels).set(duration_ms);
73    }
74}
75
76impl<S> Layer<S> for TimingMetricsLayer
77where
78    S: Subscriber + for<'a> LookupSpan<'a>,
79{
80    fn on_new_span(
81        &self,
82        attrs: &tracing::span::Attributes<'_>,
83        id: &Id,
84        ctx: tracing_subscriber::layer::Context<'_, S>,
85    ) {
86        if let Some(span) = ctx.span(id) {
87            let metadata = span.metadata();
88            let name = metadata.name();
89
90            // Only track spans at INFO level or higher to match metrics_span behavior
91            if metadata.level() <= &tracing::Level::INFO {
92                // Extract all string fields from span attributes as labels
93                let mut label_visitor = LabelVisitor::default();
94                attrs.record(&mut label_visitor);
95
96                self.span_timings.insert(
97                    id.clone(),
98                    SpanTiming {
99                        name: name.to_string(),
100                        start_time: Instant::now(),
101                        labels: label_visitor.labels,
102                    },
103                );
104            }
105        }
106    }
107
108    fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
109        // Check if this is a return event in an instrumented function
110        let mut visitor = ReturnValueVisitor { has_return: false };
111        event.record(&mut visitor);
112
113        if visitor.has_return {
114            // Get the current span
115            if let Some(span) = ctx.event_span(event) {
116                let span_id = span.id();
117
118                // Emit metric for the span that's returning
119                if let Some((_, timing)) = self.span_timings.remove(&span_id) {
120                    let duration_ms = timing.start_time.elapsed().as_millis() as f64;
121                    Self::emit_metric(&timing.name, duration_ms, &timing.labels);
122                }
123            }
124        }
125    }
126
127    fn on_close(&self, id: Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
128        // Clean up any spans that weren't emitted via return events
129        // This handles spans that don't have instrumented return values
130        if let Some((_, timing)) = self.span_timings.remove(&id) {
131            let duration_ms = timing.start_time.elapsed().as_millis() as f64;
132            Self::emit_metric(&timing.name, duration_ms, &timing.labels);
133        }
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use tracing::instrument;
140    use tracing_subscriber::{layer::SubscriberExt, Registry};
141
142    use super::*;
143
144    #[instrument(level = "info")]
145    fn example_function() -> i32 {
146        std::thread::sleep(std::time::Duration::from_millis(10));
147        42
148    }
149
150    #[test]
151    fn test_metrics_layer() {
152        let subscriber = Registry::default().with(TimingMetricsLayer::new());
153
154        tracing::subscriber::with_default(subscriber, || {
155            let result = example_function();
156            assert_eq!(result, 42);
157        });
158    }
159}
openvm_stark_sdk/metrics_tracing.rs

openvm_stark_sdk/
metrics_tracing.rs
