use crate::{
ClockDomain,
MonotonicClock,
MonotonicInstant,
Timer,
TokioRuntimeError,
TokioTimer,
};
use std::sync::Arc;
use tokio::runtime::Handle;
use tokio::time::Instant;
#[inline]
fn within_runtime<R>(runtime: &Handle, operation: impl FnOnce() -> R) -> R {
let is_current =
Handle::try_current().is_ok_and(|current| current.id() == runtime.id());
if is_current {
return operation();
}
let _runtime_guard = runtime.enter();
operation()
}
#[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
#[derive(Debug)]
pub struct TokioMonotonicClock {
domain: ClockDomain,
origin: Instant,
runtime: Handle,
}
impl TokioMonotonicClock {
#[must_use]
#[inline]
pub fn from_handle(runtime: Handle) -> Self {
let domain = ClockDomain::new();
let origin = within_runtime(&runtime, Instant::now);
Self {
domain,
origin,
runtime,
}
}
#[must_use]
#[track_caller]
#[inline]
pub fn current() -> Self {
Self::try_current().unwrap_or_else(|error| {
panic!("cannot create Tokio monotonic clock: {error}")
})
}
#[inline]
pub fn try_current() -> Result<Self, TokioRuntimeError> {
let runtime = Handle::try_current()
.map_err(|source| TokioRuntimeError::NotEntered { source })?;
Ok(Self::from_handle(runtime))
}
#[must_use]
#[inline]
pub(crate) fn same_domain_handle(&self) -> Self {
Self {
domain: self.domain,
origin: self.origin,
runtime: self.runtime.clone(),
}
}
#[must_use]
#[inline(always)]
pub(crate) const fn origin(&self) -> Instant {
self.origin
}
#[inline(always)]
pub(crate) const fn domain(&self) -> ClockDomain {
self.domain
}
#[inline]
pub(crate) fn with_runtime<R>(&self, operation: impl FnOnce() -> R) -> R {
within_runtime(&self.runtime, operation)
}
}
impl MonotonicClock for TokioMonotonicClock {
#[inline]
fn now(&self) -> MonotonicInstant {
let elapsed = self.with_runtime(|| self.origin.elapsed());
MonotonicInstant::new(self.domain, elapsed)
}
#[inline]
fn new_timer(&self) -> Arc<dyn Timer> {
Arc::new(TokioTimer::from_clock(self))
}
}