timeout-tracing 0.1.0

Returns stack and span trace when the future hits a timeout
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

#![doc = include_str!("../README.md")]

use std::{
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};

use pin_project_lite::pin_project;
use tracing::{Level, span};

use crate::waker::{TracingTimeoutWaker, TracingTimeoutWakerInner};

pub use crate::{trace::CaptureTrace, trace::DefaultTraceCapturer, trace::StackAndSpanTrace};

#[cfg(test)]
mod tests;
mod trace;
mod waker;

pub fn timeout<C, Fut>(duration: Duration, capture: C, fut: Fut) -> TimeoutFuture<C, Fut> {
    let deadline = tokio::time::sleep(duration);
    TimeoutFuture {
        deadline,
        capture: Some(capture),
        inner: fut,
    }
}

pin_project! {
    pub struct TimeoutFuture<C, Fut> {
        #[pin]
        deadline: tokio::time::Sleep,
        capture: Option<C>,
        #[pin]
        inner: Fut,
    }
}

impl<C, Fut> Future for TimeoutFuture<C, Fut>
where
    C: CaptureTrace + Send + 'static,
    Fut: Future,
{
    type Output = Result<Fut::Output, TimeoutElapsed<C::Trace>>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        // A span just so that nested timeouts had some
        let deadline_span = span!(Level::TRACE, "deadline");
        let guard = deadline_span.enter();
        match this.deadline.poll(cx) {
            Poll::Ready(()) => {
                drop(guard);

                // We hit the timeout. Do one final poll for the inner future, but collect the traces this time.
                // TODO: we don't have to call cx.waker.clone(), but it probably does not matter much since we already hit timeout
                let Some(capture) = this.capture.take() else {
                    return Poll::Ready(Err(TimeoutElapsed {
                        active_traces: Vec::new(),
                    }));
                };
                let waker_inner = TracingTimeoutWakerInner::new(capture, cx.waker().clone());
                let waker = TracingTimeoutWaker::new(waker_inner.clone()).as_std_waker();
                let mut cx2 = Context::from_waker(&waker);
                match this.inner.poll(&mut cx2) {
                    Poll::Pending => {}
                    Poll::Ready(result) => return Poll::Ready(Ok(result)),
                }
                let active_traces: Vec<_> = waker_inner.traces();
                return Poll::Ready(Err(TimeoutElapsed { active_traces }));
            }
            Poll::Pending => {}
        }
        drop(guard);
        match this.inner.poll(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(result) => Poll::Ready(Ok(result)),
        }
    }
}

#[derive(Debug)]
pub struct TimeoutElapsed<Trace> {
    pub active_traces: Vec<Trace>,
}