saluki_common/task/
instrument.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use pin_project::pin_project;
8use saluki_metrics::{static_metrics, Counter, Histogram};
9
10#[static_metrics(prefix = runtime_task, labels(task_name))]
11#[derive(Clone)]
12struct Telemetry {
13    #[metric(level = debug)]
14    poll_count: Counter,
15    #[metric(level = trace)]
16    poll_duration_seconds: Histogram,
17}
18
19/// Helper trait for instrumenting futures that are run as asynchronous tasks.
20pub trait TaskInstrument {
21    /// Instruments the future, tracking task-specific metrics about its execution.
22    ///
23    /// Whenever the resulting future is polled, two internal metrics are updated: `runtime_task.poll_count` is
24    /// incremented by one, and `runtime_task.poll_duration_seconds` records the duration of the poll operation, in
25    /// seconds. Both metrics are tagged with the task name provided here (as `task_name:<task name>`).
26    ///
27    /// In general, a unique task name should be provided where possible. If multiple tasks share the same task name,
28    /// they will all update the same metric, which will simply influence the resulting percentiles and make it more
29    /// difficult to isolate outlier poll durations.
30    fn with_task_instrumentation(self, task_name: String) -> InstrumentedTask<Self>
31    where
32        Self: Sized;
33}
34
35impl<F> TaskInstrument for F
36where
37    F: Future + Send + 'static,
38{
39    fn with_task_instrumentation(self, task_name: String) -> InstrumentedTask<Self> {
40        InstrumentedTask {
41            telemetry: Telemetry::new(task_name),
42            inner: self,
43        }
44    }
45}
46
47/// An instrumented task future.
48///
49/// This wraps a `Future` and emits telemetry about the duration of each poll operation.
50#[pin_project]
51pub struct InstrumentedTask<F> {
52    telemetry: Telemetry,
53
54    #[pin]
55    inner: F,
56}
57
58impl<F> Future for InstrumentedTask<F>
59where
60    F: Future,
61{
62    type Output = F::Output;
63
64    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
65        let this = self.project();
66
67        let poll_start = std::time::Instant::now();
68        let result = this.inner.poll(cx);
69        let poll_duration = poll_start.elapsed();
70
71        this.telemetry.poll_count.increment(1);
72        this.telemetry.poll_duration_seconds.record(poll_duration.as_secs_f64());
73
74        result
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use std::{
81        future::Future,
82        pin::Pin,
83        sync::Arc,
84        task::{Context, Poll, Wake, Waker},
85    };
86
87    use saluki_metrics::test::TestRecorder;
88
89    use super::*;
90
91    /// A future that returns `Pending` a fixed number of times before completing, letting a test
92    /// drive a known number of polls.
93    struct PendsThenReady {
94        pending_polls: usize,
95    }
96
97    impl Future for PendsThenReady {
98        type Output = ();
99
100        fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
101            if self.pending_polls == 0 {
102                Poll::Ready(())
103            } else {
104                self.pending_polls -= 1;
105                Poll::Pending
106            }
107        }
108    }
109
110    struct NoopWaker;
111
112    impl Wake for NoopWaker {
113        fn wake(self: Arc<Self>) {}
114    }
115
116    #[test]
117    fn poll_records_one_duration_sample_per_poll() {
118        let recorder = TestRecorder::default();
119        let _guard = metrics::set_default_local_recorder(&recorder);
120
121        // Two `Pending` polls followed by one `Ready` poll: three polls total.
122        let task = PendsThenReady { pending_polls: 2 }.with_task_instrumentation("poll_duration_test".to_string());
123        let mut task = Box::pin(task);
124
125        let waker = Waker::from(Arc::new(NoopWaker));
126        let mut cx = Context::from_waker(&waker);
127
128        let mut polls = 0;
129        loop {
130            polls += 1;
131            if task.as_mut().poll(&mut cx).is_ready() {
132                break;
133            }
134        }
135        assert_eq!(polls, 3);
136
137        // The documented behavior is that every poll records one `poll_duration_seconds` sample,
138        // tagged with the task name provided to `with_task_instrumentation`.
139        let samples = recorder
140            .histogram((
141                Telemetry::poll_duration_seconds_name(),
142                &[("task_name", "poll_duration_test")],
143            ))
144            .expect("poll-duration histogram should be registered");
145        assert_eq!(samples.len(), 3, "each poll must record one duration sample");
146        assert!(
147            samples.iter().all(|&sample| sample >= 0.0),
148            "recorded poll durations must be non-negative"
149        );
150
151        // Every poll also increments `poll_count`, tagged with the same task name, so after three
152        // polls the counter reads three.
153        assert_eq!(
154            recorder.counter((Telemetry::poll_count_name(), &[("task_name", "poll_duration_test")])),
155            Some(3)
156        );
157    }
158}
saluki_common/task/instrument.rs

saluki_common/task/
instrument.rs
