use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
pub type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send + 'static>>;
pub trait Spawn: Send + Sync {
fn spawn(&self, fut: BoxFuture<()>);
}
pub trait Clock: Send + Sync {
fn now_ms(&self) -> i64;
}
pub trait Sleep: Send + Sync {
fn sleep_until(&self, at_ms: i64) -> BoxFuture<()>;
}
#[derive(Clone)]
pub struct Shell {
pub spawn: Arc<dyn Spawn>,
pub clock: Arc<dyn Clock>,
pub sleep: Arc<dyn Sleep>,
}
impl std::fmt::Debug for Shell {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Shell")
.field("now_ms", &self.clock.now_ms())
.finish_non_exhaustive()
}
}
impl Shell {
pub fn tokio() -> Self {
let clock: Arc<dyn Clock> = Arc::new(SystemClock);
Shell {
spawn: Arc::new(TokioSpawn),
sleep: Arc::new(TokioSleep {
clock: clock.clone(),
}),
clock,
}
}
pub fn virtual_time(start_ms: i64) -> (Self, VirtualClock) {
let clock = VirtualClock::new(start_ms);
let shell = Shell {
spawn: Arc::new(TokioSpawn),
clock: Arc::new(clock.clone()),
sleep: Arc::new(clock.clone()),
};
(shell, clock)
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct TokioSpawn;
impl Spawn for TokioSpawn {
fn spawn(&self, fut: BoxFuture<()>) {
tokio::spawn(fut);
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct SystemClock;
impl Clock for SystemClock {
fn now_ms(&self) -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX))
.unwrap_or(0)
}
}
#[derive(Clone)]
pub struct TokioSleep {
clock: Arc<dyn Clock>,
}
impl TokioSleep {
pub fn new(clock: Arc<dyn Clock>) -> Self {
TokioSleep { clock }
}
}
impl std::fmt::Debug for TokioSleep {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TokioSleep").finish_non_exhaustive()
}
}
impl Sleep for TokioSleep {
fn sleep_until(&self, at_ms: i64) -> BoxFuture<()> {
let wait = at_ms.saturating_sub(self.clock.now_ms()).max(0);
Box::pin(tokio::time::sleep(Duration::from_millis(
u64::try_from(wait).unwrap_or(0),
)))
}
}
#[derive(Default)]
struct VirtualInner {
now: i64,
waiters: Vec<(i64, Waker)>,
}
#[derive(Clone, Default)]
pub struct VirtualClock {
inner: Arc<Mutex<VirtualInner>>,
}
impl std::fmt::Debug for VirtualClock {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("VirtualClock")
.field("now", &self.now_ms())
.finish()
}
}
impl VirtualClock {
pub fn new(start_ms: i64) -> Self {
VirtualClock {
inner: Arc::new(Mutex::new(VirtualInner {
now: start_ms,
waiters: Vec::new(),
})),
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, VirtualInner> {
self.inner
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
pub fn advance(&self, d: Duration) {
let ms = i64::try_from(d.as_millis()).unwrap_or(i64::MAX);
let now = self.now_ms().saturating_add(ms);
self.set(now);
}
pub fn set(&self, at_ms: i64) {
let due: Vec<Waker> = {
let mut g = self.lock();
g.now = g.now.max(at_ms);
let now = g.now;
let (ready, wait): (Vec<_>, Vec<_>) = g.waiters.drain(..).partition(|(t, _)| *t <= now);
g.waiters = wait;
ready.into_iter().map(|(_, w)| w).collect()
};
for w in due {
w.wake();
}
}
}
impl Clock for VirtualClock {
fn now_ms(&self) -> i64 {
self.lock().now
}
}
impl Sleep for VirtualClock {
fn sleep_until(&self, at_ms: i64) -> BoxFuture<()> {
Box::pin(VirtualSleep {
clock: self.clone(),
at: at_ms,
})
}
}
struct VirtualSleep {
clock: VirtualClock,
at: i64,
}
impl Future for VirtualSleep {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let mut g = self.clock.lock();
if g.now >= self.at {
Poll::Ready(())
} else {
g.waiters.push((self.at, cx.waker().clone()));
Poll::Pending
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn virtual_clock_moves_only_when_advanced() {
let c = VirtualClock::new(1_000);
assert_eq!(c.now_ms(), 1_000);
c.advance(Duration::from_millis(250));
assert_eq!(c.now_ms(), 1_250);
c.set(100);
assert_eq!(c.now_ms(), 1_250, "never backwards");
}
#[tokio::test]
async fn virtual_sleep_completes_after_advance() {
let c = VirtualClock::new(0);
let s = c.sleep_until(50);
let h = tokio::spawn(s);
tokio::task::yield_now().await;
assert!(!h.is_finished());
c.advance(Duration::from_millis(50));
h.await.unwrap();
}
}