agnostic_lite/tokio/
sleep.rs1use core::{future::Future, pin::Pin, task::{Context, Poll}, time::Duration};
2
3use crate::time::{AsyncLocalSleep, AsyncLocalSleepExt};
4
5pin_project_lite::pin_project! {
6 #[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}