[go: up one dir, main page]

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