async_instrumenter/
lib.rs

1pub use log;
2
3use std::{
4    future::Future,
5    pin::Pin,
6    task::{Context, Poll},
7    time::{Duration, Instant},
8};
9
10use pin_project::pin_project;
11
12// The result of a finished [`InstrumentFuture`]
13#[derive(Debug)]
14pub struct InstrumentFutureResult<R> {
15    pub result: R,
16    pub elapsed: Duration,
17}
18
19/// Wraps a future and determines exactly how long it took to execute
20///
21/// ```rust
22/// let my_fut: impl Future<Output = ()> = foobar();
23/// let res = InstrumentFuture::new(my_fut).await;
24///
25/// // get the result of `my_fut`
26/// res.result;
27///
28/// // print the elapsed time of `my_fut`
29/// println!("my_fut took {:?}", res.elapsed);
30/// ```
31#[derive(Debug)]
32#[pin_project]
33pub struct InstrumentFuture<F: Future> {
34    #[pin]
35    future: F,
36    timer: Option<Instant>,
37}
38
39impl<F: Future> InstrumentFuture<F> {
40    pub fn new(future: F) -> Self {
41        Self {
42            future,
43            timer: None,
44        }
45    }
46}
47
48impl<F: Future> Future for InstrumentFuture<F> {
49    type Output = InstrumentFutureResult<<F as Future>::Output>;
50
51    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
52        let this = self.project();
53
54        if this.timer.is_none() {
55            *this.timer = Some(Instant::now());
56        }
57
58        match this.future.poll(cx) {
59            Poll::Ready(r) => Poll::Ready(InstrumentFutureResult {
60                result: r,
61                // SAFETY: `timer` is always `Some(T)` since we set timer to Some up above when it's None
62                elapsed: unsafe { this.timer.unwrap_unchecked() }.elapsed(),
63            }),
64
65            Poll::Pending => Poll::Pending,
66        }
67    }
68}
69
70/// Debug log how long a future took to execute
71///
72/// When `debug_assertions` are enabled, does a log::debug!() with the file!(), line!(), and elapsed time.
73/// Without `debug_assertions`, simply is a no-op which executes future normally.
74///
75/// The argument to the function must be a unexecuted future. It will return a future you must `await` on.
76/// This allows you to use the created instrumenting future later if desired since it doesn't `await` immediately.
77///
78/// There is also an optional one with a custom log message. `elapsed` is provided as a keyword arg to the literal,
79/// so you must use it somewhere in there.
80///
81/// If you need custom behavior, you can make a custom instrumenting future using [`InstrumentFuture`]
82///
83/// Examples:
84///
85/// ```rust
86/// let my_fut: impl Future<Output = ()> = foobar();
87/// dbg_instrument!(my_fut).await;
88///
89/// let f = 0;
90/// dbg_instrument!("custom_log_message {f}: {elapsed:?}").await;
91/// ```
92#[macro_export]
93macro_rules! dbg_instrument {
94    ($fut:expr) => {
95        async {
96            if cfg!(debug_assertions) {
97                $crate::instrument!($fut).await
98            } else {
99                $fut.await
100            }
101        }
102    };
103
104    ($log:literal, $fut:expr) => {{
105        async {
106            if cfg!(debug_assertions) {
107                $crate::instrument!($log, $fut).await
108            } else {
109                $fut.await
110            }
111        }
112    }};
113}
114
115/// Debug log how long a future took to execute
116///
117/// As opposed to [`dbg_instrument!`], this always logs regardless whether `debug_assertions` are enabled or not
118///
119/// The argument to the function must be a unexecuted future. It will return a future you must `await` on.
120/// This allows you to use the created instrumenting future later if desired since it doesn't `await` immediately.
121///
122/// There is also an optional one with a custom log message. `elapsed` is provided as a keyword arg to the literal,
123/// so you must use it somewhere in there.
124///
125/// If you need custom behavior, you can make a custom instrumenting future using [`InstrumentFuture`]
126///
127/// Examples:
128///
129/// ```rust
130/// let my_fut: impl Future<Output = ()> = foobar();
131/// instrument!(my_fut).await;
132///
133/// let f = 0;
134/// instrument!("custom_log_message {f}: {elapsed:?}").await;
135/// ```
136#[macro_export]
137macro_rules! instrument {
138    ($fut:expr) => {
139        async { $crate::_instrument!($fut) }
140    };
141
142    ($log:literal, $fut:expr) => {{
143        async { $crate::_instrument!($log, $fut) }
144    }};
145}
146
147#[doc(hidden)]
148#[macro_export]
149macro_rules! _instrument {
150    ($fut:expr) => {{
151        let timed = $crate::InstrumentFuture::new($fut).await;
152
153        let _file = file!();
154        let _line = line!();
155        let _elapsed = timed.elapsed;
156
157        $crate::log::debug!("{_file}:{_line} completed in {_elapsed:?}");
158
159        timed.result
160    }};
161
162    ($log:literal, $fut:expr) => {{
163        let timed = $crate::InstrumentFuture::new($fut).await;
164
165        $crate::log::debug!($log, elapsed = timed.elapsed);
166
167        timed.result
168    }};
169}
async_instrumenter/lib.rs

async_instrumenter/
lib.rs
