agnostic_lite/tokio/
sleep.rs

1use core::{future::Future, pin::Pin, task::{Context, Poll}, time::Duration};
2
3use crate::time::{AsyncLocalSleep, AsyncLocalSleepExt};
4
5pin_project_lite::pin_project! {
6  /// The [`AsyncSleep`] implementation for tokio runtime
7  #[cfg_attr(docsrs, doc(cfg(all(feature = "std", feature = "tokio"))))]
8  #[repr(transparent)]
9  pub struct TokioSleep {
10    #[pin]
11    inner: ::tokio::time::Sleep,
12  }
13}
14
15impl From<::tokio::time::Sleep> for TokioSleep {
16  fn from(sleep: ::tokio::time::Sleep) -> Self {
17    Self { inner: sleep }
18  }
19}
20
21impl From<TokioSleep> for ::tokio::time::Sleep {
22  fn from(sleep: TokioSleep) -> Self {
23    sleep.inner
24  }
25}
26
27impl Future for TokioSleep {
28  type Output = ::tokio::time::Instant;
29
30  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
31    let this = self.project();
32    let ddl = this.inner.deadline();
33    match this.inner.poll(cx) {
34      Poll::Ready(_) => Poll::Ready(ddl),
35      Poll::Pending => Poll::Pending,
36    }
37  }
38}
39
40impl AsyncLocalSleep for TokioSleep {
41  type Instant = ::tokio::time::Instant;
42
43  fn reset(self: Pin<&mut Self>, deadline: Self::Instant) {
44    self.project().inner.as_mut().reset(deadline)
45  }
46}
47
48impl AsyncLocalSleepExt for TokioSleep {
49  fn sleep_local(after: Duration) -> Self
50  where
51    Self: Sized,
52  {
53    Self {
54      inner: tokio::time::sleep(after),
55    }
56  }
57
58  fn sleep_local_until(deadline: Self::Instant) -> Self
59  where
60    Self: Sized,
61  {
62    Self {
63      inner: tokio::time::sleep_until(deadline),
64    }
65  }
66}
67
68#[cfg(test)]
69mod tests {
70  use super::TokioSleep;
71  use crate::time::{AsyncSleep, AsyncSleepExt};
72  use tokio::time::{Duration, Instant};
73
74  const ORIGINAL: Duration = Duration::from_secs(1);
75  const RESET: Duration = Duration::from_secs(2);
76  const BOUND: Duration = Duration::from_millis(10);
77
78  #[tokio::test]
79  async fn test_tokio_sleep() {
80    let start = Instant::now();
81    let sleep = TokioSleep::sleep(ORIGINAL);
82    let ins = sleep.await;
83    assert!(ins >= start + ORIGINAL);
84    let elapsed = start.elapsed();
85    assert!(elapsed >= ORIGINAL && elapsed < ORIGINAL + BOUND);
86  }
87
88  #[tokio::test]
89  async fn test_tokio_sleep_until() {
90    let start = Instant::now();
91    let sleep = TokioSleep::sleep_until(start + ORIGINAL);
92    let ins = sleep.await;
93    assert!(ins >= start + ORIGINAL);
94    let elapsed = start.elapsed();
95    assert!(elapsed >= ORIGINAL && elapsed < ORIGINAL + BOUND);
96  }
97
98  #[tokio::test]
99  async fn test_tokio_sleep_reset() {
100    let start = Instant::now();
101    let sleep = TokioSleep::sleep(ORIGINAL);
102    tokio::pin!(sleep);
103    sleep.as_mut().reset(Instant::now() + RESET);
104    let ins = sleep.await;
105    assert!(ins >= start + RESET);
106    let elapsed = start.elapsed();
107    assert!(elapsed >= RESET && elapsed < RESET + BOUND);
108  }
109
110  #[tokio::test]
111  async fn test_tokio_sleep_reset2() {
112    let start = Instant::now();
113    let sleep = TokioSleep::sleep_until(start + ORIGINAL);
114    tokio::pin!(sleep);
115    sleep.as_mut().reset(Instant::now() + RESET);
116    let ins = sleep.await;
117    assert!(ins >= start + RESET);
118    let elapsed = start.elapsed();
119    assert!(elapsed >= RESET && elapsed < RESET + BOUND);
120  }
121}