Skip to main content

moirai_async/timer/
delay.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4use std::task::{Context, Poll};
5use std::time::{Duration, Instant};
6
7use crate::timer::clamped_deadline;
8use crate::timer::driver::timer_driver;
9use crate::timer::registration::TimerRegistration;
10
11/// A future that completes after a specified duration
12pub struct Delay {
13    deadline: Instant,
14    registration: Option<Arc<TimerRegistration>>,
15}
16
17impl Delay {
18    /// Create a new delay that will complete after the specified duration
19    pub fn new(duration: Duration) -> Self {
20        Self {
21            deadline: clamped_deadline(Instant::now(), duration),
22            registration: None,
23        }
24    }
25
26    /// Create a delay that completes at a specific instant
27    pub fn until(deadline: Instant) -> Self {
28        Self {
29            deadline,
30            registration: None,
31        }
32    }
33
34    /// Get the deadline for this delay
35    pub fn deadline(&self) -> Instant {
36        self.deadline
37    }
38
39    /// Reset the delay to a new duration from now
40    pub fn reset(&mut self, duration: Duration) {
41        self.deadline = clamped_deadline(Instant::now(), duration);
42        if let Some(registration) = self.registration.take() {
43            timer_driver().cancel(&registration);
44            registration.wake();
45        }
46    }
47}
48
49impl Future for Delay {
50    type Output = ();
51
52    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
53        if Instant::now() >= self.deadline {
54            if let Some(registration) = self.registration.take() {
55                timer_driver().cancel(&registration);
56            }
57            Poll::Ready(())
58        } else {
59            match &self.registration {
60                Some(registration) => registration.replace_waker(cx.waker()),
61                None => {
62                    let registration = TimerRegistration::new(cx.waker().clone());
63                    timer_driver().schedule(self.deadline, Arc::clone(&registration));
64                    self.registration = Some(registration);
65                }
66            }
67            Poll::Pending
68        }
69    }
70}
71
72impl Drop for Delay {
73    fn drop(&mut self) {
74        if let Some(registration) = self.registration.take() {
75            timer_driver().cancel(&registration);
76        }
77    }
78}