tokio_postponable_delay/
lib.rs

1use std::pin::Pin;
2use std::sync::{Arc, Mutex};
3use std::task::{Context, Poll};
4
5/// Similar to `tokio::time::Delay`,
6/// but you can push back the moment when this future will resolve
7pub struct PostponableDelay {
8    delay: tokio::time::Delay,
9    target: Arc<Mutex<(std::time::Instant, bool)>>,
10}
11
12impl PostponableDelay {
13    /// Returns a future that will resolve no sooner than `instant`
14    pub fn new(instant: std::time::Instant) -> Self {
15        let target = instant.into();
16        PostponableDelay {
17            delay: tokio::time::delay_until(target),
18            target: Arc::new(Mutex::new((instant, false))),
19        }
20    }
21
22    /// Returns a handle to allow pushing back the future's resolution
23    pub fn get_handle(&self) -> PostponableDelayHandle {
24        PostponableDelayHandle {
25            target: self.target.clone(),
26        }
27    }
28
29    fn project(
30        &mut self,
31    ) -> (
32        Pin<&mut tokio::time::Delay>,
33        &Mutex<(std::time::Instant, bool)>,
34    ) {
35        (Pin::new(&mut self.delay), &self.target)
36    }
37}
38
39/// A handle to postpone a `ResettableDelay`'s resolution
40pub struct PostponableDelayHandle {
41    target: Arc<Mutex<(std::time::Instant, bool)>>,
42}
43
44/// The result of a postpone request
45#[derive(Copy, Clone, Debug, Eq, PartialEq)]
46pub enum PostponeDelayResponse {
47    /// The delay has been successfully postponed
48    Ok,
49    /// The delay can't be postponed because it has already resolved
50    AlreadyResolved,
51    /// The request would need the task to be polled earlier than it could,
52    /// and was thus refused
53    CantResolveEarlier,
54}
55
56impl PostponeDelayResponse {
57    pub fn unwrap(self) {
58        match self {
59            PostponeDelayResponse::Ok => {}
60            other => panic!("Called .unwrap() on {:?}", other),
61        }
62    }
63}
64
65impl PostponableDelayHandle {
66    /// Attempts to postopone the corresponding `PostponableDelay`'s resolution,
67    /// returns a `PostponeDelayResponse` detailing if it succeeded.
68    #[must_use]
69    pub fn postpone(&self, target: std::time::Instant) -> PostponeDelayResponse {
70        let mut guard = self.target.lock().unwrap();
71        let previous_target = guard.0;
72        if guard.1 {
73            PostponeDelayResponse::AlreadyResolved
74        } else if target < previous_target {
75            PostponeDelayResponse::CantResolveEarlier
76        } else {
77            *&mut guard.0 = target;
78            PostponeDelayResponse::Ok
79        }
80    }
81}
82
83impl std::future::Future for PostponableDelay {
84    type Output = ();
85
86    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
87        let (delay, target) = self.project();
88        match delay.poll(cx) {
89            Poll::Pending => Poll::Pending,
90            Poll::Ready(_) => {
91                let mut guard = target.lock().unwrap();
92                let target = guard.0;
93                if target <= std::time::Instant::now() {
94                    guard.1 = true;
95                    Poll::Ready(())
96                } else {
97                    std::mem::drop(guard);
98                    self.delay = tokio::time::delay_until(target.into());
99                    self.poll(cx)
100                }
101            }
102        }
103    }
104}
105
106#[cfg(test)]
107const ERROR_MARGIN: std::time::Duration = std::time::Duration::from_millis(3);
108
109#[tokio::test]
110async fn no_resets() {
111    let target = std::time::Instant::now() + 4 * ERROR_MARGIN;
112    std::thread::sleep(2 * ERROR_MARGIN);
113    PostponableDelay::new(target).await;
114    let end = std::time::Instant::now();
115    println!("{:?}", end - target);
116    assert!(target <= end);
117    assert!(end <= target + ERROR_MARGIN);
118}
119
120#[tokio::test]
121async fn with_resets() {
122    let target = std::time::Instant::now() + 4 * ERROR_MARGIN;
123    let delay = PostponableDelay::new(target);
124    let handle = delay.get_handle();
125    std::thread::sleep(2 * ERROR_MARGIN);
126    let target = std::time::Instant::now() + 4 * ERROR_MARGIN;
127    handle.postpone(target).unwrap();
128    assert_eq!(
129        handle.postpone(target - ERROR_MARGIN),
130        PostponeDelayResponse::CantResolveEarlier
131    );
132    delay.await;
133    let end = std::time::Instant::now();
134    println!("{:?}", end - target);
135    assert_eq!(handle.postpone(end), PostponeDelayResponse::AlreadyResolved);
136    assert!(target <= end);
137    assert!(end <= target + ERROR_MARGIN);
138}
tokio_postponable_delay/lib.rs

tokio_postponable_delay/
lib.rs
