timeout_tracing/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::{
4    error::Error,
5    fmt::Display,
6    pin::Pin,
7    task::{Context, Poll},
8    time::Duration,
9};
10
11use pin_project_lite::pin_project;
12use tracing::{Level, span};
13
14use crate::waker::{TracingTimeoutWaker, TracingTimeoutWakerInner};
15
16pub use crate::{
17    trace::CaptureSpanAndStackTrace, trace::CaptureSpanTrace, trace::CaptureTrace,
18    trace::StackAndSpanTrace,
19};
20
21#[cfg(test)]
22mod tests;
23mod trace;
24mod waker;
25
26/// Drive the future `fut` to completion, while limiting its run time to `duration`.
27/// If `fut` fails to finish within `furation`, returns span traces for all active
28/// await points within `fut`.
29/// Use `capture` to specify which kind of trace should be captured.
30///
31/// # Examples
32/// ```rust
33/// # use std::time::Duration;
34/// # use tokio::time::sleep;
35/// # use timeout_tracing::{CaptureSpanTrace, timeout};
36/// # use tracing::instrument;
37/// # use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
38/// # tokio::runtime::Runtime::new()
39/// #     .unwrap()
40/// #     .block_on(async {
41/// tracing_subscriber::registry()
42///     .with(tracing_error::ErrorLayer::default())
43///     .init();
44///
45/// let result = timeout(Duration::from_secs(1), CaptureSpanTrace, long_computation()).await;
46/// assert_eq!(filter_span_trace(result.err().unwrap().to_string()), "timeout elapsed at:
47/// trace 0:
48///    0: example::step_2
49///              at example.rs:123
50///    1: example::long_computation
51///              at example.rs:123
52/// ");
53///
54/// #[instrument]
55/// async fn long_computation() {
56///     step_1().await;
57///     step_2().await;
58/// }
59/// #[instrument]
60/// async fn step_1() {
61///     sleep(Duration::from_millis(10)).await;
62/// }
63/// #[instrument]
64/// async fn step_2() {
65///     sleep(Duration::from_secs(2)).await;
66/// }
67/// # fn filter_span_trace(trace: String) -> String {
68/// #     use regex::Regex;
69/// #
70/// #     let re = Regex::new(r"at .*\.rs:([0-9_]+)").unwrap();
71/// #     let trace = re.replace_all(&trace, "at example.rs:123").to_string();
72/// #     let re = Regex::new(r"[a-z0-9:_]+::([a-z]+)").unwrap();
73/// #     let trace = re.replace_all(&trace, "example::$1").to_string();
74/// #     trace
75/// # }
76/// # });
77/// ```
78pub fn timeout<C, Fut>(duration: Duration, capture: C, fut: Fut) -> TimeoutFuture<C, Fut> {
79    let deadline = tokio::time::sleep(duration);
80    TimeoutFuture {
81        deadline,
82        capture: Some(capture),
83        inner: fut,
84    }
85}
86
87pin_project! {
88    pub struct TimeoutFuture<C, Fut> {
89        #[pin]
90        deadline: tokio::time::Sleep,
91        capture: Option<C>,
92        #[pin]
93        inner: Fut,
94    }
95}
96
97impl<C, Fut> Future for TimeoutFuture<C, Fut>
98where
99    C: CaptureTrace + Send + 'static,
100    Fut: Future,
101{
102    type Output = Result<Fut::Output, TimeoutElapsed<C::Trace>>;
103
104    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
105        let this = self.project();
106        // A span just so that nested timeouts had some
107        let deadline_span = span!(Level::TRACE, "deadline");
108        let guard = deadline_span.enter();
109        match this.deadline.poll(cx) {
110            Poll::Ready(()) => {
111                drop(guard);
112
113                // We hit the timeout. Do one final poll for the inner future, but collect the traces this time.
114                // TODO: we don't have to call cx.waker.clone(), but it probably does not matter much since we already hit timeout
115                let Some(capture) = this.capture.take() else {
116                    return Poll::Ready(Err(TimeoutElapsed {
117                        active_traces: Vec::new(),
118                    }));
119                };
120                let waker_inner = TracingTimeoutWakerInner::new(capture, cx.waker().clone());
121                let waker = TracingTimeoutWaker::new(waker_inner.clone()).as_std_waker();
122                let mut cx2 = Context::from_waker(&waker);
123                match this.inner.poll(&mut cx2) {
124                    Poll::Pending => {}
125                    Poll::Ready(result) => return Poll::Ready(Ok(result)),
126                }
127                let active_traces: Vec<_> = waker_inner.traces();
128                return Poll::Ready(Err(TimeoutElapsed { active_traces }));
129            }
130            Poll::Pending => {}
131        }
132        drop(guard);
133        match this.inner.poll(cx) {
134            Poll::Pending => Poll::Pending,
135            Poll::Ready(result) => Poll::Ready(Ok(result)),
136        }
137    }
138}
139
140#[derive(Debug)]
141pub struct TimeoutElapsed<Trace> {
142    pub active_traces: Vec<Trace>,
143}
144
145impl<Trace: Display> Display for TimeoutElapsed<Trace> {
146    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147        if self.active_traces.is_empty() {
148            f.write_str("timeout elapsed")?;
149        } else {
150            f.write_str("timeout elapsed at:\n")?;
151            for (idx, trace) in self.active_traces.iter().enumerate() {
152                writeln!(f, "trace {idx}:\n{trace}")?;
153            }
154        }
155        Ok(())
156    }
157}
158
159impl<Trace> Error for TimeoutElapsed<Trace> where Trace: std::fmt::Debug + std::fmt::Display {}
timeout_tracing/lib.rs

timeout_tracing/
lib.rs
