use std::future::Future;
use std::pin::Pin;
use std::time::Duration;
#[cfg(any(test, feature = "testing"))]
use tokio::sync::watch;
use tokio::time::Instant;
pub type Sleep = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
pub trait Clock: Send + Sync + std::fmt::Debug {
fn now(&self) -> Instant;
fn sleep_until(&self, deadline: Instant) -> Sleep;
}
#[derive(Debug, Default, Clone, Copy)]
pub struct SystemClock;
impl Clock for SystemClock {
fn now(&self) -> Instant {
Instant::now()
}
fn sleep_until(&self, deadline: Instant) -> Sleep {
Box::pin(tokio::time::sleep_until(deadline))
}
}
#[derive(Clone, Copy, Debug)]
pub struct Stopwatch {
started: Instant,
}
impl Stopwatch {
#[inline]
#[must_use]
pub fn start<C: Clock + ?Sized>(clock: &C) -> Self {
Self {
started: clock.now(),
}
}
#[inline]
#[must_use]
pub fn elapsed<C: Clock + ?Sized>(&self, clock: &C) -> Duration {
clock.now().saturating_duration_since(self.started)
}
}
#[cfg(any(test, feature = "testing"))]
#[doc(hidden)]
#[derive(Debug)]
pub struct TestClock {
base: Instant,
offset_nanos: watch::Sender<u64>,
slept_until: std::sync::Mutex<Vec<Instant>>,
}
#[cfg(any(test, feature = "testing"))]
impl TestClock {
#[must_use]
pub fn frozen() -> std::sync::Arc<TestClock> {
std::sync::Arc::new(TestClock {
base: Instant::now(),
offset_nanos: watch::Sender::new(0),
slept_until: std::sync::Mutex::new(Vec::new()),
})
}
#[must_use]
pub fn has_slept_until(&self, deadline: Instant) -> bool {
self.slept_until
.lock()
.expect("test clock lock")
.contains(&deadline)
}
pub fn advance(&self, by: Duration) {
let nanos = u64::try_from(by.as_nanos()).expect("test clock advance fits in u64 nanos");
self.offset_nanos.send_modify(|n| *n += nanos);
}
pub fn advance_stepped(&self, total: Duration, step: Duration, mut between: impl FnMut()) {
assert!(!step.is_zero(), "advance_stepped needs a non-zero step");
let mut moved = Duration::ZERO;
while moved < total {
let chunk = step.min(total - moved);
self.advance(chunk);
moved += chunk;
between();
}
}
}
#[cfg(any(test, feature = "testing"))]
impl Clock for TestClock {
fn now(&self) -> Instant {
self.base + Duration::from_nanos(*self.offset_nanos.borrow())
}
fn sleep_until(&self, deadline: Instant) -> Sleep {
self.slept_until
.lock()
.expect("test clock lock")
.push(deadline);
let base = self.base;
let mut rx = self.offset_nanos.subscribe();
Box::pin(async move {
loop {
if base + Duration::from_nanos(*rx.borrow_and_update()) >= deadline {
return;
}
if rx.changed().await.is_err() {
std::future::pending::<()>().await;
}
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_tokio_stopwatch_reports_the_time_its_clock_moved() {
let clock = TestClock::frozen();
let watch = Stopwatch::start(&*clock);
assert_eq!(watch.elapsed(&*clock), Duration::ZERO);
clock.advance(Duration::from_millis(250));
assert_eq!(watch.elapsed(&*clock), Duration::from_millis(250));
}
}