diff --git a/core/time/src/clock.rs b/core/time/src/clock.rs index fa6527f4229..08430e48ee0 100644 --- a/core/time/src/clock.rs +++ b/core/time/src/clock.rs @@ -29,19 +29,36 @@ enum ClockInner { /// it has to be replaced with a fake double, if we want our /// tests to be deterministic. /// -/// TODO: add tests. +/// # Examples +/// +/// ``` +/// use near_time::{Clock, FakeClock}; +/// +/// // In production code, use real clock +/// let clock = Clock::real(); +/// +/// // In tests, use fake clock +/// #[cfg(test)] +/// { +/// let fake_clock = FakeClock::default(); +/// let clock = fake_clock.clock(); +/// } +/// ``` #[derive(Clone)] pub struct Clock(ClockInner); impl Clock { - /// Constructor of the real clock. Use it in production code. - /// Preferably construct it directly in the main() function, + /// Creates a new instance of Clock that uses the real system clock. + /// Use this in production code, preferably constructing it directly in the main() function, /// so that it can be faked out in every other function. pub fn real() -> Clock { Clock(ClockInner::Real) } - /// Current time according to the monotone clock. + /// Returns the current time according to the monotonic clock. + /// + /// The monotonic clock is guaranteed to be always increasing and is not affected + /// by system time changes. This should be used for measuring elapsed time and timeouts. pub fn now(&self) -> Instant { match &self.0 { ClockInner::Real => Instant::now(), @@ -49,7 +66,10 @@ impl Clock { } } - /// Current time according to the system/walltime clock. + /// Returns the current UTC time according to the system clock. + /// + /// This clock can be affected by system time changes and should be used + /// when wall-clock time is needed (e.g., for timestamps in logs). pub fn now_utc(&self) -> Utc { match &self.0 { ClockInner::Real => Utc::now_utc(), @@ -57,7 +77,10 @@ impl Clock { } } - /// Cancellable. + /// Suspends the current task until the specified deadline is reached. + /// + /// If the deadline is `Infinite`, the task will be suspended indefinitely. + /// The operation is cancellable - if the future is dropped, the sleep will be cancelled. pub async fn sleep_until_deadline(&self, t: Deadline) { match t { Deadline::Infinite => std::future::pending().await, @@ -65,7 +88,9 @@ impl Clock { } } - /// Cancellable. + /// Suspends the current task until the specified instant is reached. + /// + /// The operation is cancellable - if the future is dropped, the sleep will be cancelled. pub async fn sleep_until(&self, t: Instant) { match &self.0 { ClockInner::Real => tokio::time::sleep_until(t.into()).await, @@ -73,7 +98,9 @@ impl Clock { } } - /// Cancellable. + /// Suspends the current task for the specified duration. + /// + /// The operation is cancellable - if the future is dropped, the sleep will be cancelled. pub async fn sleep(&self, d: Duration) { match &self.0 { ClockInner::Real => tokio::time::sleep(d.try_into().unwrap()).await, @@ -134,9 +161,12 @@ impl FakeClockInner { } self.instant += d; self.utc += d; + + // Wake up any waiters that have reached their deadline while let Some(earliest_waiter) = self.waiters.peek() { if earliest_waiter.deadline <= self.instant { - self.waiters.pop().unwrap().waker.send(()).ok(); + let waiter = self.waiters.pop().unwrap(); + waiter.waker.send(()).unwrap(); } else { break; } @@ -185,13 +215,16 @@ impl FakeClock { if d <= Duration::ZERO { return; } + let receiver = { let mut inner = self.0.lock().unwrap(); + let deadline = inner.now() + d; let (sender, receiver) = tokio::sync::oneshot::channel(); - let waiter = ClockWaiterInHeap { waker: sender, deadline: inner.now() + d }; + let waiter = ClockWaiterInHeap { waker: sender, deadline }; inner.waiters.push(waiter); receiver }; + receiver.await.unwrap(); } @@ -202,11 +235,13 @@ impl FakeClock { if inner.now() >= t { return; } + let (sender, receiver) = tokio::sync::oneshot::channel(); let waiter = ClockWaiterInHeap { waker: sender, deadline: t }; inner.waiters.push(waiter); receiver }; + receiver.await.unwrap(); } @@ -261,3 +296,102 @@ impl Interval { )); } } + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration as StdDuration; + + #[test] + fn test_real_clock() { + let clock = Clock::real(); + let start = clock.now(); + std::thread::sleep(StdDuration::from_millis(10)); + let end = clock.now(); + assert!(end > start); + } + + #[tokio::test] + async fn test_fake_clock_sleep() { + let fake = FakeClock::default(); + let clock = fake.clock(); + let start = clock.now(); + + // Create a task that sleeps + let sleep_task = tokio::spawn({ + let clock = clock.clone(); + async move { + println!("Sleep task starting at {:?}", clock.now()); + clock.sleep(Duration::seconds(5)).await; + println!("Sleep task woke up at {:?}", clock.now()); + clock.now() + } + }); + + // Give the sleep task a chance to start and register its waiter + tokio::task::yield_now().await; + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + + println!("Advancing clock by 3 seconds from {:?}", fake.now()); + fake.advance(Duration::seconds(3)); + println!("Clock is now at {:?}", fake.now()); + + // Sleep task should still be waiting + assert!(!sleep_task.is_finished()); + + println!("Advancing clock by 3 more seconds"); + fake.advance(Duration::seconds(3)); + println!("Clock is now at {:?}", fake.now()); + + // Now sleep task should complete with a timeout to prevent hanging + let end = tokio::time::timeout(std::time::Duration::from_secs(1), sleep_task) + .await + .expect("sleep_task timed out") + .expect("sleep_task panicked"); + + assert_eq!(end.signed_duration_since(start), Duration::seconds(6)); + } + + #[tokio::test] + async fn test_fake_clock_sleep_until() { + let fake = FakeClock::default(); + let clock = fake.clock(); + let start = clock.now(); + let wake_time = start + Duration::seconds(10); + + // Create a task that sleeps until specific time + let sleep_task = tokio::spawn({ + let clock = clock.clone(); + async move { + clock.sleep_until(wake_time).await; + clock.now() + } + }); + + // Advance clock to just before wake time + fake.advance_until(wake_time - Duration::seconds(1)); + + // Sleep task should still be waiting + assert!(!sleep_task.is_finished()); + + // Advance clock past wake time + fake.advance_until(wake_time + Duration::seconds(1)); + + // Now sleep task should complete + let end = sleep_task.await.unwrap(); + assert!(end >= wake_time); + } + + #[test] + fn test_fake_clock_utc() { + let fake = FakeClock::default(); + let clock = fake.clock(); + let start_utc = clock.now_utc(); + + // Advance clock by 1 hour + fake.advance(Duration::hours(1)); + + let end_utc = clock.now_utc(); + assert_eq!(end_utc - start_utc, Duration::hours(1)); + } +}