use std::cell::RefCell;
use std::sync::{Arc, OnceLock};
use std::time::{Duration as StdDuration, Instant};
use async_trait::async_trait;
use harn_clock::Clock;
use time::OffsetDateTime;
thread_local! {
static MOCK_CLOCK_STACK: RefCell<Vec<Arc<MockClock>>> = const { RefCell::new(Vec::new()) };
}
fn process_start() -> &'static Instant {
static PROCESS_START: OnceLock<Instant> = OnceLock::new();
PROCESS_START.get_or_init(Instant::now)
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct ClockInstant(StdDuration);
impl ClockInstant {
pub fn duration_since(self, earlier: Self) -> StdDuration {
self.0.saturating_sub(earlier.0)
}
pub fn as_millis(self) -> u128 {
self.0.as_millis()
}
}
pub struct ClockOverrideGuard;
impl Drop for ClockOverrideGuard {
fn drop(&mut self) {
let _ = MOCK_CLOCK_STACK.try_with(|slot| {
slot.borrow_mut().pop();
});
}
}
#[derive(Debug)]
pub struct MockClock {
inner: Arc<harn_clock::PausedClock>,
}
impl MockClock {
pub fn new(now: OffsetDateTime) -> Arc<Self> {
Arc::new(Self {
inner: harn_clock::PausedClock::new(now),
})
}
pub fn at_wall_ms(wall_ms: i64) -> Arc<Self> {
let nanos = (wall_ms as i128).saturating_mul(1_000_000);
let now =
OffsetDateTime::from_unix_timestamp_nanos(nanos).unwrap_or(OffsetDateTime::UNIX_EPOCH);
Self::new(now)
}
pub fn monotonic_now(&self) -> ClockInstant {
ClockInstant(StdDuration::from_millis(
self.inner.monotonic_ms().max(0) as u64
))
}
pub fn now_wall_ms(&self) -> i64 {
self.now_utc().unix_timestamp_nanos() as i64 / 1_000_000
}
pub fn now_monotonic_ms(&self) -> i64 {
self.inner.monotonic_ms()
}
pub fn set_sync(&self, now: OffsetDateTime) {
self.inner.set(now);
}
pub fn advance_std_sync(&self, duration: StdDuration) {
self.inner.advance(duration);
}
pub async fn set(&self, now: OffsetDateTime) {
self.inner.set(now);
}
pub async fn advance(&self, duration: time::Duration) {
self.inner.advance_time(duration);
}
pub async fn advance_std(&self, duration: StdDuration) {
self.inner.advance(duration);
}
pub async fn advance_ticks(&self, ticks: u32, tick: StdDuration) {
self.inner.advance_ticks(ticks, tick);
}
}
#[async_trait]
impl Clock for MockClock {
fn now_utc(&self) -> OffsetDateTime {
self.inner.now_utc()
}
fn monotonic_ms(&self) -> i64 {
self.inner.monotonic_ms()
}
async fn sleep(&self, duration: StdDuration) {
self.inner.sleep(duration).await;
}
async fn sleep_until_utc(&self, deadline: OffsetDateTime) {
self.inner.sleep_until_utc(deadline).await;
}
}
pub fn install_override(clock: Arc<MockClock>) -> ClockOverrideGuard {
MOCK_CLOCK_STACK.with(|slot| {
slot.borrow_mut().push(clock);
});
ClockOverrideGuard
}
pub fn active_mock_clock() -> Option<Arc<MockClock>> {
MOCK_CLOCK_STACK.with(|slot| slot.borrow().last().cloned())
}
pub fn is_mocked() -> bool {
MOCK_CLOCK_STACK.with(|slot| !slot.borrow().is_empty())
}
pub fn clear_overrides() {
MOCK_CLOCK_STACK.with(|slot| {
slot.borrow_mut().clear();
});
}
pub fn now_utc() -> OffsetDateTime {
active_mock_clock()
.map(|clock| clock.now_utc())
.unwrap_or_else(OffsetDateTime::now_utc)
}
pub fn now_ms() -> i64 {
now_utc().unix_timestamp_nanos() as i64 / 1_000_000
}
pub fn instant_now() -> ClockInstant {
active_mock_clock()
.map(|clock| clock.monotonic_now())
.unwrap_or_else(|| ClockInstant(process_start().elapsed()))
}
pub fn advance(duration: StdDuration) {
if let Some(clock) = active_mock_clock() {
clock.advance_std_sync(duration);
}
}
pub async fn sleep(duration: StdDuration) {
if duration.is_zero() {
return;
}
if let Some(mock) = active_mock_clock() {
mock.sleep(duration).await;
return;
}
tokio::time::sleep(duration).await;
}