Skip to main content

es_entity/clock/
controller.rs

1use chrono::{DateTime, Utc};
2
3use std::{sync::Arc, time::Duration};
4
5use super::manual::ManualClock;
6
7/// Controller for manual time operations.
8///
9/// This is only available for manual clocks and provides methods to
10/// advance time, set time, and inspect pending wake events.
11///
12/// Created alongside a [`ClockHandle`](crate::ClockHandle) via
13/// [`ClockHandle::manual()`](crate::ClockHandle::manual).
14#[derive(Clone)]
15pub struct ClockController {
16    pub(crate) clock: Arc<ManualClock>,
17}
18
19impl ClockController {
20    /// Advance time by the given duration.
21    ///
22    /// Wake events are processed in chronological order. If you advance by
23    /// 1 day and there are sleeps scheduled at 1 hour and 2 hours, they will
24    /// wake at their scheduled times (seeing the correct `now()` value),
25    /// not at +1 day.
26    ///
27    /// Returns the number of wake events that were processed.
28    ///
29    /// # Example
30    ///
31    /// ```rust
32    /// use es_entity::clock::ClockHandle;
33    /// use std::time::Duration;
34    ///
35    /// # async fn example() {
36    /// let (clock, ctrl) = ClockHandle::manual();
37    /// let t0 = clock.now();
38    ///
39    /// let clock2 = clock.clone();
40    /// let handle = tokio::spawn(async move {
41    ///     clock2.sleep(Duration::from_secs(3600)).await;
42    ///     clock2.now() // Will be t0 + 1 hour
43    /// });
44    ///
45    /// tokio::task::yield_now().await; // Let task register its sleep
46    ///
47    /// // Advance 1 day - task wakes at exactly +1 hour
48    /// ctrl.advance(Duration::from_secs(86400)).await;
49    ///
50    /// let wake_time = handle.await.unwrap();
51    /// assert_eq!(wake_time, t0 + chrono::Duration::hours(1));
52    /// # }
53    /// ```
54    pub async fn advance(&self, duration: Duration) -> usize {
55        self.clock.advance(duration).await
56    }
57
58    /// Advance to the next pending wake event.
59    ///
60    /// Returns the time that was advanced to, or `None` if there are no
61    /// pending wake events.
62    ///
63    /// This is useful for step-by-step testing where you want to process
64    /// events one at a time.
65    ///
66    /// # Example
67    ///
68    /// ```rust
69    /// use es_entity::clock::ClockHandle;
70    /// use std::time::Duration;
71    ///
72    /// # async fn example() {
73    /// let (clock, ctrl) = ClockHandle::manual();
74    /// let t0 = clock.now();
75    ///
76    /// // Spawn tasks with different sleep durations
77    /// let c = clock.clone();
78    /// tokio::spawn(async move { c.sleep(Duration::from_secs(60)).await; });
79    /// let c = clock.clone();
80    /// tokio::spawn(async move { c.sleep(Duration::from_secs(120)).await; });
81    ///
82    /// tokio::task::yield_now().await;
83    ///
84    /// // Step through one wake at a time
85    /// let t1 = ctrl.advance_to_next_wake().await;
86    /// assert_eq!(t1, Some(t0 + chrono::Duration::seconds(60)));
87    ///
88    /// let t2 = ctrl.advance_to_next_wake().await;
89    /// assert_eq!(t2, Some(t0 + chrono::Duration::seconds(120)));
90    ///
91    /// let t3 = ctrl.advance_to_next_wake().await;
92    /// assert_eq!(t3, None); // No more pending wakes
93    /// # }
94    /// ```
95    pub async fn advance_to_next_wake(&self) -> Option<DateTime<Utc>> {
96        self.clock.advance_to_next_wake().await
97    }
98
99    /// Set the time to a specific value.
100    ///
101    /// **Warning**: Unlike `advance()`, this does NOT process wake events in order.
102    /// All tasks whose wake time has passed will see the new time when they wake.
103    /// Use this for "jump ahead, don't care about intermediate events" scenarios.
104    ///
105    /// For deterministic testing, prefer `advance()` or `advance_to_next_wake()`.
106    pub fn set_time(&self, time: DateTime<Utc>) {
107        self.clock.set_time(time);
108        // Wake all tasks that are now past their wake time
109        self.clock.wake_tasks_at(time.timestamp_millis());
110    }
111
112    /// Get the number of pending wake events.
113    ///
114    /// This is useful for testing to verify that tasks have registered
115    /// their sleeps before advancing time.
116    pub fn pending_wake_count(&self) -> usize {
117        self.clock.pending_wake_count()
118    }
119
120    /// Get the current time.
121    ///
122    /// This is equivalent to calling `now()` on the associated `ClockHandle`.
123    pub fn now(&self) -> DateTime<Utc> {
124        self.clock.now()
125    }
126
127    /// Clear all pending wake events.
128    pub fn clear_pending_wakes(&self) {
129        self.clock.clear_pending_wakes();
130    }
131
132    /// Reset clock to a specific time and clear all pending wakes.
133    ///
134    /// Useful for test isolation between test cases.
135    pub fn reset_to(&self, time: DateTime<Utc>) {
136        self.clock.set_time(time);
137        self.clock.clear_pending_wakes();
138    }
139}
140
141impl std::fmt::Debug for ClockController {
142    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
143        f.debug_struct("ClockController")
144            .field("now", &self.clock.now())
145            .field("pending_wakes", &self.clock.pending_wake_count())
146            .finish()
147    }
148}