use std::{
future::Future,
pin::Pin,
sync::{
atomic::{AtomicI64, Ordering},
Arc,
},
};
use anyhow::Error;
use chrono::{DateTime, TimeDelta, TimeZone as _, Timelike as _, Utc};
use tokio::{select, sync::Notify, time::sleep};
use super::{CheatClockFactory, CheatClockRemote, Clock, ClockBase, SAFETY_REFRESH_DELAY};
trait StableCheatClockBase {
fn time_s(&self) -> &Arc<AtomicI64>;
fn stable_now(&self) -> DateTime<Utc> {
let timestamp = self.time_s().load(Ordering::Relaxed);
timestamp_to_date(timestamp)
}
}
#[derive(Clone, Debug, Default)]
pub struct StableCheatClock {
time_s: Arc<AtomicI64>,
remote_time_s: Arc<AtomicI64>,
notify: Arc<Notify>,
}
impl StableCheatClockBase for StableCheatClock {
fn time_s(&self) -> &Arc<AtomicI64> {
&self.time_s
}
}
impl ClockBase for StableCheatClock {
fn now(&self) -> DateTime<Utc> {
self.stable_now()
}
}
impl Clock for StableCheatClock {
fn sleep(&self, delay: TimeDelta) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
self.sleep_until(self.now() + delay)
}
fn sleep_until(
&self,
deadline: DateTime<Utc>,
) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
Box::pin(async move {
let time_deadline = deadline.timestamp();
loop {
let time = self.time_s.load(Ordering::Relaxed);
let remote_time = self.remote_time_s.load(Ordering::Relaxed);
#[cfg(any(test, debug_assertions))]
assert!(time <= remote_time);
if time != remote_time && time != time_deadline {
if time_deadline >= remote_time {
self.time_s.store(remote_time, Ordering::Relaxed);
self.notify.notify_waiters();
} else {
self.time_s.store(time_deadline, Ordering::Relaxed);
}
}
if self.now().timestamp() >= time_deadline {
break;
}
select! {
_ = sleep(SAFETY_REFRESH_DELAY) => {},
_ = self.notify.notified() => {}
}
}
})
}
fn clone_box(&self) -> Box<dyn Clock> {
Box::new(self.clone())
}
fn synchronize(&self) {
self.time_s.store(
self.remote_time_s.load(Ordering::Relaxed),
Ordering::Relaxed,
);
}
}
#[derive(Clone, Debug)]
pub struct StableCheatClockRemote {
time_s: Arc<AtomicI64>,
remote_time_s: Arc<AtomicI64>,
notify: Arc<Notify>,
}
impl StableCheatClockBase for StableCheatClockRemote {
fn time_s(&self) -> &Arc<AtomicI64> {
&self.time_s
}
}
impl ClockBase for StableCheatClockRemote {
fn now(&self) -> DateTime<Utc> {
self.stable_now()
}
}
impl CheatClockRemote for StableCheatClockRemote {
#[cfg(test)]
fn not_blocking(&self) {
self.remote_time_s.store(i64::MAX, Ordering::Relaxed);
}
fn advance_toward(&self, date: DateTime<Utc>) {
self.remote_time_s
.store(date.timestamp(), Ordering::Relaxed);
self.notify.notify_waiters();
}
fn remote_now(&self) -> DateTime<Utc> {
let timestamp = self.remote_time_s.load(Ordering::Relaxed);
timestamp_to_date(timestamp)
}
fn notify(&self) -> &Notify {
&self.notify
}
fn is_sync(&self) -> bool {
self.time_s.load(Ordering::Relaxed) == self.remote_time_s.load(Ordering::Relaxed)
}
fn clone_box(&self) -> Box<dyn CheatClockRemote> {
Box::new(self.clone())
}
}
pub struct StableCheatClockFactory;
impl CheatClockFactory for StableCheatClockFactory {
fn get_clock(
&self,
date: DateTime<Utc>,
) -> Result<(Box<dyn CheatClockRemote>, Box<dyn Clock>), anyhow::Error> {
let now = Utc::now();
if now < date {
return Err(Error::msg("Can't set time in the future"));
}
let time_s = Arc::new(AtomicI64::new(date.timestamp()));
let remote_time_s = Arc::new(AtomicI64::new(date.timestamp()));
let notify = Arc::new(Notify::new());
Ok((
Box::new(StableCheatClockRemote {
time_s: time_s.clone(),
remote_time_s: remote_time_s.clone(),
notify: notify.clone(),
}),
Box::new(StableCheatClock {
time_s,
remote_time_s,
notify,
}),
))
}
}
fn timestamp_to_date(timestamp: i64) -> DateTime<Utc> {
Utc.timestamp_opt(timestamp, 0)
.single()
.unwrap_or_else(|| -> DateTime<Utc> {
tracing::error!("Failed to convert timestamp {} to DateTime<Utc>", timestamp);
Utc::now()
})
.with_nanosecond(0)
.unwrap_or_else(|| -> DateTime<Utc> {
tracing::error!("Failed get rid of nanoseconds");
Utc::now()
})
}
#[cfg(test)]
mod tests {
use chrono::{DateTime, TimeDelta, Utc};
use std::{
io::{self, Write},
str::FromStr as _,
sync::{
atomic::{AtomicI16, Ordering},
Arc,
},
};
use tokio::time::sleep as tokio_sleep;
use super::{super::REMOTE_REFRESH_RATE, CheatClockFactory as _, StableCheatClockFactory};
const TOLERANCE: TimeDelta = TimeDelta::milliseconds(10);
const DELAY: TimeDelta = TimeDelta::seconds(1);
#[tokio::test]
async fn sleep() {
let date: DateTime<Utc> = DateTime::from_str("2012-12-12 00:00:00Z").unwrap();
let (remote, clock) = StableCheatClockFactory {}.get_clock(date).unwrap();
remote.not_blocking();
assert_eq!(clock.now(), date);
let before = Utc::now();
clock.sleep(DELAY).await;
let after = Utc::now();
assert!(after - before < TOLERANCE);
assert_eq!(clock.now(), date + DELAY);
}
#[tokio::test]
async fn sleep_until() {
let date: DateTime<Utc> = DateTime::from_str("2012-12-12 00:00:00Z").unwrap();
let (remote, clock) = StableCheatClockFactory {}.get_clock(date).unwrap();
remote.not_blocking();
assert_eq!(clock.now(), date);
let before = Utc::now();
clock.sleep_until(date + DELAY).await;
let after = Utc::now();
assert!(after - before < TOLERANCE);
assert_eq!(clock.now(), date + DELAY);
}
#[tokio::test]
async fn remote() {
let date = DateTime::from_str("2012-12-12 00:00:00Z").unwrap();
let (remote, clock) = StableCheatClockFactory {}.get_clock(date).unwrap();
assert_eq!(clock.now(), date);
assert_eq!(remote.now(), date);
let marker = Arc::new(AtomicI16::new(0));
let marker_clone = marker.clone();
tokio::spawn(async move {
clock.sleep(DELAY).await;
marker_clone.store(1, Ordering::Relaxed);
clock.sleep(DELAY + DELAY).await;
marker_clone.store(2, Ordering::Relaxed);
clock
.sleep_until(DateTime::from_str("2020-01-01 00:00:00Z").unwrap())
.await;
panic!("should not reach here");
});
assert_eq!(marker.load(Ordering::Relaxed), 0);
tokio_sleep(DELAY.to_std().unwrap()).await;
tokio_sleep(REMOTE_REFRESH_RATE + TOLERANCE.to_std().unwrap()).await;
assert_eq!(marker.load(Ordering::Relaxed), 0);
remote.advance_toward(date + DELAY + DELAY);
println!("Loop running");
io::stdout().flush().unwrap();
tokio_sleep(TOLERANCE.to_std().unwrap()).await;
assert_eq!(marker.load(Ordering::Relaxed), 1);
}
}