relay_server/utils/sleep_handle.rs
1use std::future::Future;
2use std::pin::Pin;
3use std::task::Poll;
4use std::time::Duration;
5
6/// A future wrapper around [`tokio::time::Sleep`].
7///
8/// When initialized with [`SleepHandle::idle`], this future is pending indefinitely every time it
9/// is polled. To initiate a delay, use [`set`](Self::set). After the delay has passed, the future
10/// resolves with `()` **exactly once** and resets to idle. To reset the future while it is
11/// sleeping, use [`reset`](Self::reset).
12#[derive(Debug)]
13pub struct SleepHandle(Option<Pin<Box<tokio::time::Sleep>>>);
14
15impl SleepHandle {
16 /// Creates [`SleepHandle`] and sets its internal state to an indefinitely pending future.
17 pub fn idle() -> Self {
18 Self(None)
19 }
20
21 /// Resets the internal state to an indefinitely pending future.
22 pub fn reset(&mut self) {
23 self.0 = None;
24 }
25
26 /// Sets the internal state to a future that will yield after `duration` time has elapsed.
27 pub fn set(&mut self, duration: Duration) {
28 self.0 = Some(Box::pin(tokio::time::sleep(duration)));
29 }
30
31 /// Checks wether the internal state is currently pending indefinite.
32 pub fn is_idle(&self) -> bool {
33 self.0.is_none()
34 }
35}
36
37impl Future for SleepHandle {
38 type Output = ();
39
40 fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
41 let poll = match &mut self.0 {
42 Some(sleep) => Pin::new(sleep).poll(cx),
43 None => Poll::Pending,
44 };
45
46 if poll.is_ready() {
47 self.reset();
48 }
49
50 poll
51 }
52}