timeout_tracing/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::{
4    pin::Pin,
5    task::{Context, Poll},
6    time::Duration,
7};
8
9use pin_project_lite::pin_project;
10use tracing::{Level, span};
11
12use crate::waker::{TracingTimeoutWaker, TracingTimeoutWakerInner};
13
14pub use crate::{trace::CaptureTrace, trace::DefaultTraceCapturer, trace::StackAndSpanTrace};
15
16#[cfg(test)]
17mod tests;
18mod trace;
19mod waker;
20
21pub fn timeout<C, Fut>(duration: Duration, capture: C, fut: Fut) -> TimeoutFuture<C, Fut> {
22    let deadline = tokio::time::sleep(duration);
23    TimeoutFuture {
24        deadline,
25        capture: Some(capture),
26        inner: fut,
27    }
28}
29
30pin_project! {
31    pub struct TimeoutFuture<C, Fut> {
32        #[pin]
33        deadline: tokio::time::Sleep,
34        capture: Option<C>,
35        #[pin]
36        inner: Fut,
37    }
38}
39
40impl<C, Fut> Future for TimeoutFuture<C, Fut>
41where
42    C: CaptureTrace + Send + 'static,
43    Fut: Future,
44{
45    type Output = Result<Fut::Output, TimeoutElapsed<C::Trace>>;
46
47    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
48        let this = self.project();
49        // A span just so that nested timeouts had some
50        let deadline_span = span!(Level::TRACE, "deadline");
51        let guard = deadline_span.enter();
52        match this.deadline.poll(cx) {
53            Poll::Ready(()) => {
54                drop(guard);
55
56                // We hit the timeout. Do one final poll for the inner future, but collect the traces this time.
57                // TODO: we don't have to call cx.waker.clone(), but it probably does not matter much since we already hit timeout
58                let Some(capture) = this.capture.take() else {
59                    return Poll::Ready(Err(TimeoutElapsed {
60                        active_traces: Vec::new(),
61                    }));
62                };
63                let waker_inner = TracingTimeoutWakerInner::new(capture, cx.waker().clone());
64                let waker = TracingTimeoutWaker::new(waker_inner.clone()).as_std_waker();
65                let mut cx2 = Context::from_waker(&waker);
66                match this.inner.poll(&mut cx2) {
67                    Poll::Pending => {}
68                    Poll::Ready(result) => return Poll::Ready(Ok(result)),
69                }
70                let active_traces: Vec<_> = waker_inner.traces();
71                return Poll::Ready(Err(TimeoutElapsed { active_traces }));
72            }
73            Poll::Pending => {}
74        }
75        drop(guard);
76        match this.inner.poll(cx) {
77            Poll::Pending => Poll::Pending,
78            Poll::Ready(result) => Poll::Ready(Ok(result)),
79        }
80    }
81}
82
83#[derive(Debug)]
84pub struct TimeoutElapsed<Trace> {
85    pub active_traces: Vec<Trace>,
86}
timeout_tracing/lib.rs

timeout_tracing/
lib.rs
