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}