io2/
time.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5    time::{Duration, Instant},
6};
7
8use crate::executor::CURRENT_TASK_CONTEXT;
9
10#[must_use = "futures do nothing unless you `.await` or poll them"]
11pub struct NotifyWhen {
12    timer: Option<Instant>,
13}
14
15impl Future for NotifyWhen {
16    type Output = ();
17
18    fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
19        let fut = self.get_mut();
20        match fut.timer.take() {
21            Some(timer) => {
22                CURRENT_TASK_CONTEXT.with_borrow_mut(|ctx| {
23                    let ctx = ctx.as_mut().unwrap();
24                    ctx.notify_when(timer);
25                });
26                Poll::Pending
27            }
28            None => Poll::Ready(()),
29        }
30    }
31}
32
33pub fn sleep(duration: Duration) -> NotifyWhen {
34    let now = Instant::now();
35    let timer = now.checked_add(duration).unwrap();
36    NotifyWhen { timer: Some(timer) }
37}
38
39pub fn sleep_until(instant: Instant) -> NotifyWhen {
40    NotifyWhen {
41        timer: Some(instant),
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use crate::executor::ExecutorConfig;
48
49    use super::*;
50
51    #[test]
52    #[ignore]
53    fn test_sleep() {
54        ExecutorConfig::new()
55            .run(async {
56                for _ in 0..20 {
57                    sleep(Duration::from_secs(1)).await;
58                    println!("SLEEPING");
59                }
60            })
61            .unwrap();
62    }
63}