use crate::timer::internal::tokio_runtime_liveness::TokioRuntimeLiveness;
use crate::timer::internal::tokio_runtime_liveness_registry::TokioRuntimeLivenessRegistry;
use crate::timer::internal::tokio_timer_future::TokioTimerFuture;
use crate::{
MonotonicClock,
MonotonicInstant,
TimeError,
Timer,
TimerFuture,
TimerUnavailableError,
TokioMonotonicClock,
TokioRuntimeError,
};
#[cfg(coverage)]
use std::sync::atomic::{
AtomicBool,
Ordering,
};
use std::{
panic::{
AssertUnwindSafe,
catch_unwind,
},
sync::{
Arc,
OnceLock,
},
time::Duration,
};
use tokio::runtime::Handle;
use tokio::time::Instant;
#[cfg(coverage)]
static PANIC_NEXT_SLEEP_POLL: AtomicBool = AtomicBool::new(false);
#[cfg(coverage)]
pub fn panic_next_tokio_timer_sleep_poll() {
PANIC_NEXT_SLEEP_POLL.store(true, Ordering::Release);
}
#[cfg(coverage)]
pub(crate) fn take_tokio_timer_sleep_poll_panic() -> bool {
PANIC_NEXT_SLEEP_POLL.swap(false, Ordering::AcqRel)
}
#[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
#[derive(Debug)]
pub struct TokioTimer {
clock: TokioMonotonicClock,
liveness: OnceLock<Arc<TokioRuntimeLiveness>>,
}
impl TokioTimer {
#[must_use]
#[inline]
pub fn from_handle(runtime: Handle) -> Self {
Self {
clock: TokioMonotonicClock::from_handle(runtime),
liveness: OnceLock::new(),
}
}
#[must_use]
#[track_caller]
#[inline]
pub fn current() -> Self {
Self::try_current().unwrap_or_else(|error| {
panic!("cannot create Tokio timer: {error}")
})
}
#[inline]
pub fn try_current() -> Result<Self, TokioRuntimeError> {
TokioMonotonicClock::try_current().map(|clock| Self {
clock,
liveness: OnceLock::new(),
})
}
#[must_use]
#[inline]
pub fn from_clock(clock: &TokioMonotonicClock) -> Self {
Self {
clock: clock.same_domain_handle(),
liveness: OnceLock::new(),
}
}
fn native_deadline(
&self,
deadline: MonotonicInstant,
) -> Result<Instant, TimeError> {
deadline.validate_domain(self.clock.domain())?;
self.clock
.origin()
.checked_add(deadline.elapsed_since_origin())
.ok_or(TimeError::InstantOverflow)
}
fn runtime_liveness(&self) -> Arc<TokioRuntimeLiveness> {
if let Some(liveness) = self.liveness.get() {
return Arc::clone(liveness);
}
let liveness = TokioRuntimeLivenessRegistry::current();
let _ = self.liveness.set(liveness);
Arc::clone(
self.liveness
.get()
.expect("Tokio timer liveness should be initialized"),
)
}
fn schedule(
&self,
deadline: Instant,
now: Instant,
) -> Result<TimerFuture, TimeError> {
if deadline <= now {
return Ok(Box::pin(std::future::ready(Ok(()))));
}
let sleep = catch_unwind(AssertUnwindSafe(|| {
tokio::time::sleep_until(deadline)
}))
.map_err(|_| TimeError::TimerUnavailable {
source: TimerUnavailableError::TimeDriverDisabled,
})?;
let liveness = self.runtime_liveness();
Ok(Box::pin(TokioTimerFuture::new(sleep, liveness)))
}
}
impl Timer for TokioTimer {
#[inline(always)]
fn clock(&self) -> &dyn MonotonicClock {
&self.clock
}
fn at(&self, deadline: MonotonicInstant) -> Result<TimerFuture, TimeError> {
let deadline = self.native_deadline(deadline)?;
self.clock
.with_runtime(|| self.schedule(deadline, Instant::now()))
}
#[inline]
fn after(&self, duration: Duration) -> Result<TimerFuture, TimeError> {
self.clock.with_runtime(|| {
let now = Instant::now();
let deadline = now
.checked_add(duration)
.ok_or(TimeError::InstantOverflow)?;
self.schedule(deadline, now)
})
}
}