fuel_core_metrics/futures/
metered_future.rs

1use crate::futures::{
2    future_tracker::FutureTracker,
3    FuturesMetrics,
4};
5use std::{
6    future::Future,
7    pin::Pin,
8    task::{
9        Context,
10        Poll,
11    },
12};
13
14pin_project_lite::pin_project! {
15    /// Future that tracks its execution with [`FutureTracker`] and reports it back when done
16    /// to [`FuturesMetrics`].
17    #[derive(Debug, Clone)]
18    #[must_use = "futures do nothing unless you `.await` or poll them"]
19    pub struct MeteredFuture<F> {
20        #[pin]
21        future: FutureTracker<F>,
22        metrics: FuturesMetrics,
23    }
24}
25
26impl<F> MeteredFuture<F> {
27    /// Create a new `MeteredFuture` with the given future and metrics.
28    pub fn new(future: F, metrics: FuturesMetrics) -> Self {
29        Self {
30            future: FutureTracker::new(future),
31            metrics,
32        }
33    }
34}
35
36impl<F> Future for MeteredFuture<F>
37where
38    F: Future,
39{
40    type Output = F::Output;
41
42    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
43        let this = self.project();
44
45        match this.future.poll(cx) {
46            Poll::Ready(output) => {
47                let output = output.extract(this.metrics);
48                Poll::Ready(output)
49            }
50            Poll::Pending => Poll::Pending,
51        }
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use std::time::Duration;
59
60    #[tokio::test]
61    async fn hybrid_time_correct() {
62        let future = async {
63            tokio::time::sleep(Duration::from_secs(2)).await;
64            std::thread::sleep(Duration::from_secs(1));
65        };
66        let metrics = FuturesMetrics::obtain_futures_metrics("test");
67        let wrapper_future = MeteredFuture::new(future, metrics.clone());
68        let _ = wrapper_future.await;
69        let busy = Duration::from_nanos(metrics.busy.get());
70        let idle = Duration::from_nanos(metrics.idle.get());
71        assert_eq!(idle.as_secs(), 2);
72        assert_eq!(busy.as_secs(), 1);
73    }
74}