compio_runtime/time/
future.rs

1use std::{
2    cell::RefCell,
3    pin::Pin,
4    rc::Rc,
5    task::{Context, Poll},
6    time::{Duration, Instant},
7};
8
9use pin_project_lite::pin_project;
10
11use crate::{
12    Runtime,
13    time::{Elapsed, TimerRuntime, runtime::TimerKey, sleep_until},
14};
15
16#[derive(Debug)]
17pub(crate) struct TimerFuture {
18    key: TimerKey,
19    rt: Rc<RefCell<TimerRuntime>>,
20}
21
22impl TimerFuture {
23    /// Try to create a new `TimerFuture` if the instant is in the future;
24    /// otherwise, a `None` will be returned.
25    ///
26    /// # Panics
27    ///
28    /// Panic if not running under a `Runtime`.
29    pub fn try_new(instant: Instant) -> Option<Self> {
30        Runtime::with_current(|rt| {
31            let key = rt.timer_runtime.borrow_mut().insert(instant)?;
32            Some(Self {
33                key,
34                rt: rt.timer_runtime.clone(),
35            })
36        })
37    }
38}
39
40impl Future for TimerFuture {
41    type Output = ();
42
43    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
44        self.rt.borrow_mut().poll_timer(cx, &self.key)
45    }
46}
47
48impl Drop for TimerFuture {
49    fn drop(&mut self) {
50        self.rt.borrow_mut().cancel(&self.key)
51    }
52}
53
54compio_driver::assert_not_impl!(TimerFuture, Send);
55compio_driver::assert_not_impl!(TimerFuture, Sync);
56
57/// Future returned by [`sleep`](super::sleep) and
58/// [`sleep_until`](super::sleep_until).
59#[must_use = "Futures do nothing unless polled."]
60#[derive(Debug)]
61pub struct Sleep(Option<TimerFuture>);
62
63impl Sleep {
64    #[inline]
65    pub(crate) fn new(instant: Instant) -> Self {
66        Sleep(TimerFuture::try_new(instant))
67    }
68}
69
70impl Future for Sleep {
71    type Output = ();
72
73    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
74        if let Some(timer) = self.0.as_mut() {
75            Pin::new(timer).poll(cx)
76        } else {
77            Poll::Ready(())
78        }
79    }
80}
81
82pin_project! {
83    /// Future returned by [`timeout`](timeout) and [`timeout_at`](timeout_at).
84    #[must_use = "Futures do nothing unless polled."]
85    #[derive(Debug)]
86    pub struct Timeout<F> {
87        #[pin]
88        fut: F,
89        sleep: Sleep,
90    }
91}
92
93impl<F: Future> Timeout<F> {
94    pub(crate) fn new(instant: Instant, fut: F) -> Self {
95        Self {
96            fut,
97            sleep: Sleep::new(instant),
98        }
99    }
100}
101
102impl<F: Future> Future for Timeout<F> {
103    type Output = Result<F::Output, Elapsed>;
104
105    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
106        let me = self.project();
107
108        if let Poll::Ready(out) = me.fut.poll(cx) {
109            return Poll::Ready(Ok(out));
110        }
111
112        match Pin::new(me.sleep).poll(cx) {
113            Poll::Ready(()) => Poll::Ready(Err(Elapsed(()))),
114            Poll::Pending => Poll::Pending,
115        }
116    }
117}
118
119/// Interval returned by [`interval`] and [`interval_at`]
120///
121/// This type allows you to wait on a sequence of instants with a certain
122/// duration between each instant. Unlike calling [`sleep`] in a loop, this lets
123/// you count the time spent between the calls to [`sleep`] as well.
124///
125/// [`sleep`]: super::sleep
126/// [`interval`]: super::interval
127/// [`interval_at`]: super::interval_at
128#[derive(Debug)]
129pub struct Interval {
130    first_ticked: bool,
131    start: Instant,
132    period: Duration,
133}
134
135impl Interval {
136    pub(crate) fn new(start: Instant, period: Duration) -> Self {
137        Self {
138            first_ticked: false,
139            start,
140            period,
141        }
142    }
143
144    /// Completes when the next instant in the interval has been reached.
145    ///
146    /// See [`interval`](super::interval) and
147    /// [`interval_at`](super::interval_at).
148    pub async fn tick(&mut self) -> Instant {
149        if !self.first_ticked {
150            sleep_until(self.start).await;
151            self.first_ticked = true;
152            self.start
153        } else {
154            let now = Instant::now();
155            let next = now + self.period
156                - Duration::from_nanos(
157                    ((now - self.start).as_nanos() % self.period.as_nanos()) as _,
158                );
159            sleep_until(next).await;
160            next
161        }
162    }
163}
compio_runtime/time/future.rs

compio_runtime/time/
future.rs
