worker/
delay.rs

1use std::{
2    cell::Cell,
3    future::Future,
4    pin::Pin,
5    rc::Rc,
6    task::{Context, Poll},
7    time::Duration,
8};
9
10use wasm_bindgen::{prelude::Closure, JsCast};
11
12fn timeout_to_clear(awoken: bool, timeout_id: Option<i32>) -> Option<i32> {
13    if awoken {
14        None
15    } else {
16        timeout_id
17    }
18}
19
20/// A [Future] for asynchronously waiting.
21///
22/// # Example:
23/// ```rust,ignore
24/// use std::time::Duration;
25/// use worker::Delay;
26///
27/// let duration = Duration::from_millis(1000);
28///
29/// // Waits a second
30/// Delay::from(duration).await;
31/// ```
32#[derive(Debug)]
33#[pin_project::pin_project(PinnedDrop)]
34pub struct Delay {
35    inner: Duration,
36    closure: Option<Closure<dyn FnMut()>>,
37    timeout_id: Option<i32>,
38    awoken: Rc<Cell<bool>>,
39}
40
41impl Future for Delay {
42    type Output = ();
43
44    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
45        let this = self.project();
46
47        if !this.awoken.get() {
48            if this.closure.is_none() {
49                let awoken = this.awoken.clone();
50                let callback_ref = this.closure.get_or_insert_with(move || {
51                    let waker = cx.waker().clone();
52                    let wake = Box::new(move || {
53                        waker.wake_by_ref();
54                        awoken.set(true);
55                    });
56
57                    Closure::wrap_assert_unwind_safe(wake as _)
58                });
59
60                // Then get that closure back and pass it to setTimeout so we can get woken up later.
61                let global: web_sys::WorkerGlobalScope = js_sys::global().unchecked_into();
62                let timeout_id = global
63                    .set_timeout_with_callback_and_timeout_and_arguments_0(
64                        callback_ref.as_ref().unchecked_ref::<js_sys::Function>(),
65                        this.inner.as_millis() as i32,
66                    )
67                    .unwrap();
68                *this.timeout_id = Some(timeout_id);
69            }
70
71            Poll::Pending
72        } else {
73            Poll::Ready(())
74        }
75    }
76}
77
78impl From<Duration> for Delay {
79    fn from(inner: Duration) -> Self {
80        Self {
81            inner,
82            closure: None,
83            timeout_id: None,
84            awoken: Rc::new(Cell::default()),
85        }
86    }
87}
88
89/// SAFETY: If, for whatever reason, the delay is dropped before the future is ready JS will invoke
90/// a dropped future causing memory safety issues. To avoid this we will just clean up the timeout
91/// if we drop the delay, cancelling the timeout.
92#[pin_project::pinned_drop]
93impl PinnedDrop for Delay {
94    fn drop(self: Pin<&'_ mut Self>) {
95        let this = self.project();
96
97        if let Some(id) = timeout_to_clear(this.awoken.get(), *this.timeout_id) {
98            let global: web_sys::WorkerGlobalScope = js_sys::global().unchecked_into();
99            global.clear_timeout_with_handle(id);
100        }
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::timeout_to_clear;
107
108    #[test]
109    fn clears_timeout_when_not_awoken() {
110        assert_eq!(timeout_to_clear(false, Some(42)), Some(42));
111    }
112
113    #[test]
114    fn does_not_clear_timeout_when_awoken() {
115        assert_eq!(timeout_to_clear(true, Some(42)), None);
116    }
117
118    #[test]
119    fn does_not_clear_without_timeout_id() {
120        assert_eq!(timeout_to_clear(false, None), None);
121    }
122}
worker/delay.rs

worker/
delay.rs
