es_entity/clock/
sleep.rs

1use pin_project::{pin_project, pinned_drop};
2use tokio::time::Sleep;
3
4use std::{
5    future::Future,
6    pin::Pin,
7    sync::{
8        Arc,
9        atomic::{AtomicU64, Ordering},
10    },
11    task::{Context, Poll},
12    time::Duration,
13};
14
15use super::{inner::ClockInner, manual::ManualClock};
16
17/// Counter for unique sleep IDs.
18static NEXT_SLEEP_ID: AtomicU64 = AtomicU64::new(0);
19
20/// Generate a unique sleep ID.
21fn next_sleep_id() -> u64 {
22    NEXT_SLEEP_ID.fetch_add(1, Ordering::Relaxed)
23}
24
25/// A future that completes after a duration has elapsed on the clock.
26///
27/// Created by [`ClockHandle::sleep`](crate::ClockHandle::sleep).
28#[pin_project(PinnedDrop)]
29pub struct ClockSleep {
30    #[pin]
31    inner: ClockSleepInner,
32}
33
34#[pin_project(project = ClockSleepInnerProj)]
35enum ClockSleepInner {
36    Realtime {
37        #[pin]
38        sleep: Sleep,
39    },
40    Manual {
41        wake_at_ms: i64,
42        sleep_id: u64,
43        clock: Arc<ManualClock>,
44        registered: bool,
45        /// If true, this wake is registered in coalesce_wakes instead of pending_wakes.
46        coalesceable: bool,
47    },
48}
49
50impl ClockSleep {
51    pub(crate) fn new(clock_inner: &ClockInner, duration: Duration) -> Self {
52        Self::new_inner(clock_inner, duration, false)
53    }
54
55    pub(crate) fn new_coalesceable(clock_inner: &ClockInner, duration: Duration) -> Self {
56        Self::new_inner(clock_inner, duration, true)
57    }
58
59    fn new_inner(clock_inner: &ClockInner, duration: Duration, coalesceable: bool) -> Self {
60        let inner = match clock_inner {
61            ClockInner::Realtime(rt) => ClockSleepInner::Realtime {
62                sleep: rt.sleep(duration),
63            },
64            ClockInner::Manual(manual) => {
65                let added = i64::try_from(duration.as_millis()).unwrap_or(i64::MAX);
66                let wake_at_ms = manual.now_ms().saturating_add(added);
67
68                ClockSleepInner::Manual {
69                    wake_at_ms,
70                    sleep_id: next_sleep_id(),
71                    clock: Arc::clone(manual),
72                    registered: false,
73                    coalesceable,
74                }
75            }
76        };
77
78        Self { inner }
79    }
80}
81
82impl Future for ClockSleep {
83    type Output = ();
84
85    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
86        let this = self.project();
87
88        match this.inner.project() {
89            ClockSleepInnerProj::Realtime { sleep } => sleep.poll(cx),
90
91            ClockSleepInnerProj::Manual {
92                wake_at_ms,
93                sleep_id,
94                clock,
95                registered,
96                coalesceable,
97            } => {
98                // Check if we've reached wake time
99                if clock.now_ms() >= *wake_at_ms {
100                    return Poll::Ready(());
101                }
102
103                // Register for wake notification if not already done
104                if !*registered {
105                    if *coalesceable {
106                        clock.register_coalesce_wake(*wake_at_ms, *sleep_id, cx.waker().clone());
107                    } else {
108                        clock.register_wake(*wake_at_ms, *sleep_id, cx.waker().clone());
109                    }
110                    *registered = true;
111                }
112
113                Poll::Pending
114            }
115        }
116    }
117}
118
119#[pinned_drop]
120impl PinnedDrop for ClockSleep {
121    fn drop(self: Pin<&mut Self>) {
122        // Clean up pending wake registration if cancelled
123        if let ClockSleepInner::Manual {
124            sleep_id,
125            clock,
126            registered: true,
127            ..
128        } = &self.inner
129        {
130            clock.cancel_wake(*sleep_id);
131        }
132    }
133}
134
135/// A future that completes with a timeout after a duration has elapsed on the clock.
136///
137/// Created by [`ClockHandle::timeout`](crate::ClockHandle::timeout).
138#[pin_project]
139pub struct ClockTimeout<F> {
140    #[pin]
141    future: F,
142    #[pin]
143    sleep: ClockSleep,
144    completed: bool,
145}
146
147impl<F> ClockTimeout<F> {
148    pub(crate) fn new(clock_inner: &ClockInner, duration: Duration, future: F) -> Self {
149        Self {
150            future,
151            sleep: ClockSleep::new(clock_inner, duration),
152            completed: false,
153        }
154    }
155}
156
157impl<F: Future> Future for ClockTimeout<F> {
158    type Output = Result<F::Output, Elapsed>;
159
160    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
161        let this = self.project();
162
163        if *this.completed {
164            panic!("ClockTimeout polled after completion");
165        }
166
167        // Check the future first
168        if let Poll::Ready(output) = this.future.poll(cx) {
169            *this.completed = true;
170            return Poll::Ready(Ok(output));
171        }
172
173        // Check if timeout elapsed
174        if let Poll::Ready(()) = this.sleep.poll(cx) {
175            *this.completed = true;
176            return Poll::Ready(Err(Elapsed));
177        }
178
179        Poll::Pending
180    }
181}
182
183/// Error returned when a timeout expires.
184#[derive(Debug, Clone, Copy, PartialEq, Eq)]
185pub struct Elapsed;
186
187impl std::fmt::Display for Elapsed {
188    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189        write!(f, "deadline has elapsed")
190    }
191}
192
193impl std::error::Error for Elapsed {}
es_entity/clock/sleep.rs

es_entity/clock/
sleep.rs
