timeout-tracing 0.1.2

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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use std::{
    sync::{Arc, Mutex},
    task::{RawWaker, RawWakerVTable, Waker},
};

use crate::trace::CaptureTrace;

pub struct TracingTimeoutWakerInner<C: CaptureTrace + Send + 'static> {
    active_traces: Mutex<Vec<Option<C::Trace>>>,
    capture: C,
    inner_waker: Waker,
}

impl<C: CaptureTrace + Send + 'static> TracingTimeoutWakerInner<C> {
    pub fn new(capture: C, inner_waker: Waker) -> Arc<Self> {
        Arc::new(Self {
            active_traces: Mutex::new(Vec::with_capacity(4)),
            capture,
            inner_waker,
        })
    }

    pub fn traces(&self) -> Vec<C::Trace> {
        std::mem::take(&mut *self.active_traces.lock().unwrap())
            .into_iter()
            .flatten()
            .collect()
    }
}

pub struct TracingTimeoutWaker<C: CaptureTrace + Send + 'static> {
    inner: Arc<TracingTimeoutWakerInner<C>>,
    idx: Option<usize>,
}

impl<C> TracingTimeoutWaker<C>
where
    C: CaptureTrace + Send + 'static,
{
    fn vtable() -> &'static RawWakerVTable {
        &RawWakerVTable::new(
            Self::raw_clone,
            Self::raw_wake,
            Self::raw_wake_by_ref,
            Self::raw_drop,
        )
    }

    pub fn new(inner: Arc<TracingTimeoutWakerInner<C>>) -> Box<Self> {
        Box::new(Self { inner, idx: None })
    }

    pub fn as_std_waker(self: Box<Self>) -> Waker {
        let data = Box::into_raw(self);
        unsafe { Waker::new(data as *const (), Self::vtable()) }
    }

    fn clone(&self) -> Box<Self> {
        let trace = self.inner.capture.capture();
        let idx = {
            let mut traces = self.inner.active_traces.lock().unwrap();
            let idx = traces.len();
            traces.push(Some(trace));
            idx
        };
        Box::new(Self {
            inner: self.inner.clone(),
            idx: Some(idx),
        })
    }

    unsafe fn raw_clone(data: *const ()) -> RawWaker {
        let this = unsafe { &*(data as *const Self) };
        let cloned = this.clone();
        RawWaker::new(Box::into_raw(cloned) as *const (), Self::vtable())
    }

    fn wake(&self) {
        self.inner.inner_waker.wake_by_ref()
    }
    unsafe fn raw_wake(data: *const ()) {
        let this = unsafe { &*(data as *const Self) };
        this.wake();
    }
    fn wake_by_ref(&self) {
        self.inner.inner_waker.wake_by_ref()
    }
    unsafe fn raw_wake_by_ref(data: *const ()) {
        let this = unsafe { &*(data as *const Self) };
        this.wake_by_ref();
    }
    unsafe fn raw_drop(data: *const ()) {
        let this = unsafe { Box::<Self>::from_raw(data as *mut Self) };
        drop(this);
    }
}

impl<C: CaptureTrace + Send + 'static> Drop for TracingTimeoutWaker<C> {
    fn drop(&mut self) {
        let mut traces = self.inner.active_traces.lock().unwrap();
        if let Some(idx) = self.idx {
            // `TimeoutFuture::poll` takes data from `traces.traces`
            if !traces.is_empty() {
                traces[idx] = None;
            }
        }
    }
}