dope_runtime/
time.rs

1use std::{
2    future::Future,
3    io,
4    pin::Pin,
5    ptr::NonNull,
6    task::{Context, Poll},
7    time::{Duration, Instant},
8};
9
10use dope_core::driver::CompletionBackend;
11
12use super::executor;
13
14pub async fn sleep_with<B: CompletionBackend>(
15    ctx: executor::RuntimeCtx<B>,
16    duration: Duration,
17) -> io::Result<()> {
18    if duration.is_zero() {
19        return Ok(());
20    }
21    Sleep::new(ctx.timers(), Instant::now() + duration).await;
22    Ok(())
23}
24
25pub async fn sleep<B: CompletionBackend>(duration: Duration) -> io::Result<()> {
26    sleep_with(executor::current::<B>(), duration).await
27}
28
29pub fn timeout_with<B, F>(
30    ctx: executor::RuntimeCtx<B>,
31    duration: Duration,
32    fut: F,
33) -> impl Future<Output = io::Result<F::Output>>
34where
35    B: CompletionBackend,
36    F: Future,
37{
38    Timeout::new(ctx.timers(), duration, fut)
39}
40
41pub fn timeout<B, F>(duration: Duration, fut: F) -> impl Future<Output = io::Result<F::Output>>
42where
43    B: CompletionBackend,
44    F: Future,
45{
46    timeout_with(executor::current::<B>(), duration, fut)
47}
48
49struct Timeout<F>
50where
51    F: Future,
52{
53    fut: F,
54    deadline: Instant,
55    sleep: Sleep,
56}
57
58impl<F> Future for Timeout<F>
59where
60    F: Future,
61{
62    type Output = io::Result<F::Output>;
63
64    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
65        let this = unsafe { self.get_unchecked_mut() };
66
67        if let Poll::Ready(v) = unsafe { Pin::new_unchecked(&mut this.fut) }.poll(cx) {
68            return Poll::Ready(Ok(v));
69        }
70
71        let now = Instant::now();
72        if now >= this.deadline {
73            return Poll::Ready(Err(io::Error::new(io::ErrorKind::TimedOut, "timeout")));
74        }
75
76        match Pin::new(&mut this.sleep).poll(cx) {
77            Poll::Pending => Poll::Pending,
78            Poll::Ready(()) => Poll::Ready(Err(io::Error::new(io::ErrorKind::TimedOut, "timeout"))),
79        }
80    }
81}
82
83impl<F> Timeout<F>
84where
85    F: Future,
86{
87    fn new(timers: NonNull<executor::TimerHeap>, duration: Duration, fut: F) -> Self {
88        let deadline = Instant::now() + duration;
89        Self {
90            fut,
91            deadline,
92            sleep: Sleep::new(timers, deadline),
93        }
94    }
95}
96
97struct Sleep {
98    deadline: Instant,
99    id: Option<u64>,
100    timers: NonNull<executor::TimerHeap>,
101}
102
103impl Sleep {
104    fn new(timers: NonNull<executor::TimerHeap>, deadline: Instant) -> Self {
105        Self {
106            deadline,
107            id: None,
108            timers,
109        }
110    }
111
112    #[inline(always)]
113    fn timers(&self) -> &executor::TimerHeap {
114        unsafe { self.timers.as_ref() }
115    }
116}
117
118impl Drop for Sleep {
119    fn drop(&mut self) {
120        if let Some(id) = self.id.take() {
121            self.timers().cancel(id);
122        }
123    }
124}
125
126impl Future for Sleep {
127    type Output = ();
128
129    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
130        let now = Instant::now();
131        if now >= self.deadline {
132            if let Some(id) = self.id.take() {
133                self.timers().cancel(id);
134            }
135            return Poll::Ready(());
136        }
137
138        let deadline = self.deadline;
139        let id = self.id;
140        let new_id = self.timers().register(id, deadline, cx.waker());
141        self.id = Some(new_id);
142        Poll::Pending
143    }
144}
dope_runtime/time.rs

dope_runtime/
time.rs
