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}