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
12/// A [Future] for asynchronously waiting.
13///
14/// # Example:
15/// ```rust,ignore
16/// use std::time::Duration;
17/// use worker::Delay;
18///
19/// let duration = Duration::from_millis(1000);
20///
21/// // Waits a second
22/// Delay::from(duration).await;
23/// ```
24#[derive(Debug)]
25#[pin_project::pin_project(PinnedDrop)]
26pub struct Delay {
27    inner: Duration,
28    closure: Option<Closure<dyn FnMut()>>,
29    timeout_id: Option<i32>,
30    awoken: Rc<Cell<bool>>,
31}
32
33impl Future for Delay {
34    type Output = ();
35
36    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
37        let this = self.project();
38
39        if !this.awoken.get() {
40            if this.closure.is_none() {
41                let awoken = this.awoken.clone();
42                let callback_ref = this.closure.get_or_insert_with(move || {
43                    let waker = cx.waker().clone();
44                    let wake = Box::new(move || {
45                        waker.wake_by_ref();
46                        awoken.set(true);
47                    });
48
49                    Closure::wrap(wake as _)
50                });
51
52                // Then get that closure back and pass it to setTimeout so we can get woken up later.
53                let global: web_sys::WorkerGlobalScope = js_sys::global().unchecked_into();
54                let timeout_id = global
55                    .set_timeout_with_callback_and_timeout_and_arguments_0(
56                        callback_ref.as_ref().unchecked_ref::<js_sys::Function>(),
57                        this.inner.as_millis() as i32,
58                    )
59                    .unwrap();
60                *this.timeout_id = Some(timeout_id);
61            }
62
63            Poll::Pending
64        } else {
65            Poll::Ready(())
66        }
67    }
68}
69
70impl From<Duration> for Delay {
71    fn from(inner: Duration) -> Self {
72        Self {
73            inner,
74            closure: None,
75            timeout_id: None,
76            awoken: Rc::new(Cell::default()),
77        }
78    }
79}
80
81/// SAFETY: If, for whatever reason, the delay is dropped before the future is ready JS will invoke
82/// a dropped future causing memory safety issues. To avoid this we will just clean up the timeout
83/// if we drop the delay, cancelling the timeout.
84#[pin_project::pinned_drop]
85impl PinnedDrop for Delay {
86    fn drop(self: Pin<&'_ mut Self>) {
87        let this = self.project();
88
89        // If we've already completed the future we don't need to clear the timeout.
90        if this.awoken.get() {
91            return;
92        }
93
94        if let Some(id) = this.timeout_id {
95            crate::console_debug!("{:#?} has been dropped", &this.inner);
96            let global: web_sys::WorkerGlobalScope = js_sys::global().unchecked_into();
97            global.clear_timeout_with_handle(*id);
98        }
99    }
100}
worker/delay.rs

worker/
delay.rs
