use std::sync::{
Arc,
atomic::{AtomicI64, Ordering},
};
use chrono::{DateTime, TimeDelta, TimeZone as _, Timelike as _, Utc};
use tokio::{select, sync::Notify, time::sleep};
use super::{Clock, ClockBase, ClockFactory, ClockRemote, SAFETY_REFRESH_DELAY};
use crate::core::{CoreError, CoreResult};
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 {
const NAME: &'static str = "StableCheatClock";
fn now(&self) -> DateTime<Utc> {
self.stable_now()
}
}
impl Clock for StableCheatClock {
type Remote = StableCheatClockRemote;
async fn sleep(&self, delay: TimeDelta) {
self.sleep_until(self.now() + delay).await
}
async fn sleep_until(&self, deadline: DateTime<Utc>) {
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 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 {
const NAME: &'static str = "StableCheatClockRemote";
fn now(&self) -> DateTime<Utc> {
self.stable_now()
}
}
impl ClockRemote 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)
}
}
pub struct StableCheatClockFactory {
pub date: DateTime<Utc>,
}
impl ClockFactory for StableCheatClockFactory {
type _Clock = StableCheatClock;
fn build(self) -> CoreResult<(Option<StableCheatClockRemote>, StableCheatClock)> {
let now = Utc::now();
if now < self.date {
return Err(CoreError::param_error("Can't set time in the future"));
}
let time_s = Arc::new(AtomicI64::new(self.date.timestamp()));
let remote_time_s = Arc::new(AtomicI64::new(self.date.timestamp()));
let notify = Arc::new(Notify::new());
Ok((
Some(StableCheatClockRemote {
time_s: time_s.clone(),
remote_time_s: remote_time_s.clone(),
notify: notify.clone(),
}),
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::{
Arc,
atomic::{AtomicI16, Ordering},
},
};
use tokio::time::sleep as tokio_sleep;
use super::{
super::REMOTE_REFRESH_RATE, Clock as _, ClockBase as _, ClockFactory as _,
ClockRemote 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 { date }.build().unwrap();
remote.unwrap().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 { date }.build().unwrap();
remote.unwrap().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 { date }.build().unwrap();
let remote = _remote.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);
}
}