use std::panic::AssertUnwindSafe;
use std::panic::catch_unwind;
use std::sync::Arc;
use std::sync::OnceLock;
#[cfg(coverage)]
use std::sync::atomic::AtomicBool;
#[cfg(coverage)]
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::runtime::Handle;
use tokio::time::Instant;
use tokio::time::sleep_until;
use crate::MonotonicClock;
use crate::MonotonicInstant;
use crate::TimeError;
use crate::Timer;
use crate::TimerFuture;
use crate::TimerUnavailableError;
use crate::TokioMonotonicClock;
use crate::TokioRuntimeError;
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;
#[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(|| 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)
})
}
}