Skip to main content

wasm_bindgen_spawn/
join.rs

1use std::any::Any;
2use std::marker::PhantomData;
3use std::pin::Pin;
4use std::task::{Context, Poll};
5
6use crate::util::{Value, ValueReceiver, ValueReceiverAsync, WorkerPanic, WorkerResult};
7
8/// Handle for joining a thread
9///
10/// This can be used as a drop-in replacement for [`std::thread::JoinHandle`]
11/// as long as you are not calling `.thread()` (which currently doesn't have a good use case -
12/// if you do have one please open an issue on GitHub).
13///
14/// `JoinHandle` also implements [`IntoFuture`](std::future::IntoFuture) so you can join a thread
15/// asynchronously by `await`-ing it.
16pub struct JoinHandle<T: Send + 'static> {
17    id: usize,
18    recv: ValueReceiver,
19    _marker: PhantomData<T>,
20}
21
22impl<T: Send + 'static> JoinHandle<T> {
23    pub(crate) fn new(id: usize, recv: ValueReceiver) -> Self {
24        Self {
25            id,
26            recv,
27            _marker: PhantomData,
28        }
29    }
30
31    /// Block the current thread until the thread is finished.
32    /// Returns the value returned by the thread's main closure/future.
33    ///
34    /// This function should behave similarly to [`std::thread::JoinHandle::join`].
35    ///
36    /// To asynchronously join the thread you can `await` the join handle instead of calling
37    /// `.join()`.
38    ///
39    /// # Note about panicking
40    /// If `panic=abort`, the panic will still be caught and this function will return
41    /// `Err` with a generic message, instead of triggering another panic.
42    ///
43    /// For more information on unwind and catching panics, see
44    /// [Working with Panic](https://wbgspawn.pistonite.dev/panic.html) in the book
45    pub fn join(self) -> Result<T, Box<dyn Any + Send + 'static>> {
46        handle_join_result(self.id, self.recv.recv())
47    }
48
49    /// Check if the thread has finished executing, or panicked.
50    /// This can be used to implement non-blocking join.
51    pub fn is_finished(&self) -> bool {
52        self.recv.has_message() || self.recv.is_closed()
53    }
54}
55
56/// [`IntoFuture`] implementation for [`JoinHandle`]
57pub struct AsyncJoinHandle<T: Send + 'static> {
58    id: usize,
59    recv: ValueReceiverAsync,
60    _marker: PhantomData<T>,
61}
62
63impl<T: Send + 'static> IntoFuture for JoinHandle<T> {
64    type Output = Result<T, Box<dyn Any + Send + 'static>>;
65    type IntoFuture = AsyncJoinHandle<T>;
66
67    fn into_future(self) -> Self::IntoFuture {
68        AsyncJoinHandle {
69            id: self.id,
70            recv: self.recv.into_future(),
71            _marker: PhantomData,
72        }
73    }
74}
75
76impl<T: Send + 'static> Future for AsyncJoinHandle<T> {
77    type Output = Result<T, Box<dyn Any + Send + 'static>>;
78
79    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
80        let id = self.id;
81        // safety: recv is pinned while self is
82        let recv = unsafe { self.map_unchecked_mut(|s| &mut s.recv) };
83        match recv.poll(cx) {
84            Poll::Ready(x) => Poll::Ready(handle_join_result(id, x)),
85            Poll::Pending => Poll::Pending,
86        }
87    }
88}
89
90fn handle_join_result<T>(
91    id: usize,
92    result: Result<WorkerResult, oneshot::RecvError>,
93) -> Result<T, Box<dyn Any + Send + 'static>> {
94    // recv() will only error if somehow the thread terminated without sending a value
95    let result = match result {
96        Ok(x) => x,
97        Err(_) => return Err(Box::new(format!("thread {id} is disconnected"))),
98    };
99    // cast the value back from void* to Box<T>
100    let value: Value = match result {
101        Ok(x) => x,
102        Err(WorkerPanic { payload: Some(e) }) => {
103            return Err(e);
104        }
105        Err(WorkerPanic { payload: None }) => {
106            if cfg!(panic = "unwind") {
107                // see https://wasm-bindgen.github.io/wasm-bindgen/reference/handling-aborts.html
108                return Err(Box::new(format!(
109                    "thread {id} encountered a non-recoverable hard abort!",
110                )));
111            }
112            return Err(Box::new(format!("thread {id} panicked or aborted!")));
113        }
114    };
115    // safety: join handle created in spawn should have the same type T
116    let value: Box<T> = unsafe { value.into_box_unchecked() };
117    Ok(*value)
118}