use std::{ future::Future, sync::{Arc, OnceLock}, time::{Duration, Instant}, }; static NOOP_WAKER: OnceLock = OnceLock::new(); fn noop_waker() -> &'static std::task::Waker { NOOP_WAKER.get_or_init(|| std::task::Waker::from(Arc::new(NoopWaker))) } struct NoopWaker; impl std::task::Wake for NoopWaker { fn wake(self: std::sync::Arc) {} fn wake_by_ref(self: &std::sync::Arc) {} } pub(crate) type LocalBoxFuture<'a, T> = std::pin::Pin + 'a>>; pub(crate) trait NowOrNever: Future { fn now_or_never(self) -> Option where Self: Sized, { let fut = std::pin::pin!(self); let mut cx = std::task::Context::from_waker(noop_waker()); match fut.poll(&mut cx) { std::task::Poll::Pending => None, std::task::Poll::Ready(out) => Some(out), } } } pub(crate) trait WithTimeout: Future { fn with_timeout(self, duration: Duration) -> tokio::time::Timeout where Self: Sized, { tokio::time::timeout(duration, self) } } pub(crate) trait WithMetrics: Future { fn with_metrics(self, name: &'static str) -> MetricsFuture where Self: Sized, { MetricsFuture { future: self, metrics: Metrics { name, start: Instant::now(), complete: false, }, } } } impl NowOrNever for F where F: Future {} impl WithMetrics for F where F: Future {} impl WithTimeout for F where F: Future {} pin_project_lite::pin_project! { pub(crate) struct MetricsFuture { #[pin] future: F, metrics: Metrics, } } struct Metrics { name: &'static str, start: Instant, complete: bool, } impl Future for MetricsFuture where F: Future, { type Output = F::Output; fn poll( self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> std::task::Poll { let this = self.project(); let out = std::task::ready!(this.future.poll(cx)); this.metrics.complete = true; std::task::Poll::Ready(out) } } impl Drop for Metrics { fn drop(&mut self) { metrics::histogram!(self.name, "complete" => self.complete.to_string()) .record(self.start.elapsed().as_secs_f64()); } }