feat(wasm): Improve wasm join handle to implement more tokio methods (#5088)

This change adds support for the abort/abort_handle/is_finished methods
onto the JoinHandle shim for Wasm targets.

Signed-off-by: Daniel Salinas

Author: zzorba

crates/matrix-sdk-common/Cargo.toml

@@ -34,17 +34,17 @@ ruma.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true
-tokio = { workspace = true, features = ["rt", "time"] }
 tracing.workspace = true
 uniffi = { workspace = true, optional = true }
 
-[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+[target.'cfg(not(target_family = "wasm"))'.dependencies]
 # Enable the test macro.
-tokio = { workspace = true, features = ["rt", "macros"] }
+tokio = { workspace = true, features = ["rt", "time", "macros"] }
 
-[target.'cfg(target_arch = "wasm32")'.dependencies]
+[target.'cfg(target_family = "wasm")'.dependencies]
 futures-util = { workspace = true, features = ["channel"] }
 gloo-timers = { workspace = true, features = ["futures"] }
+tokio = { workspace = true, features = ["macros"] }
 tracing-subscriber = { workspace = true, features = ["fmt", "ansi"] }
 wasm-bindgen.workspace = true
 wasm-bindgen-futures = { version = "0.4.33", optional = true }
@@ -57,7 +57,7 @@ matrix-sdk-test-macros = { path = "../../testing/matrix-sdk-test-macros" }
 proptest.workspace = true
 wasm-bindgen-test.workspace = true
 
-[target.'cfg(target_arch = "wasm32")'.dev-dependencies]
+[target.'cfg(target_family = "wasm")'.dev-dependencies]
 # Enable the JS feature for getrandom.
 getrandom = { workspace = true, default-features = false, features = ["js"] }
 js-sys.workspace = true

crates/matrix-sdk-common/src/executor.rs

@@ -12,88 +12,135 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! Abstraction over an executor so we can spawn tasks under WASM the same way
+//! Abstraction over an executor so we can spawn tasks under Wasm the same way
 //! we do usually.
 
-#[cfg(target_arch = "wasm32")]
-use std::{
-    future::Future,
-    pin::Pin,
-    task::{Context, Poll},
-};
-
-#[cfg(target_arch = "wasm32")]
-pub use futures_util::future::Aborted as JoinError;
-#[cfg(target_arch = "wasm32")]
-use futures_util::{
-    future::{AbortHandle, Abortable, RemoteHandle},
-    FutureExt,
-};
-#[cfg(not(target_arch = "wasm32"))]
-pub use tokio::task::{spawn, JoinError, JoinHandle};
-
-#[cfg(target_arch = "wasm32")]
-pub fn spawn<F, T>(future: F) -> JoinHandle<T>
-where
-    F: Future<Output = T> + 'static,
-{
-    let (future, remote_handle) = future.remote_handle();
-    let (abort_handle, abort_registration) = AbortHandle::new_pair();
-    let future = Abortable::new(future, abort_registration);
-
-    wasm_bindgen_futures::spawn_local(async {
-        // Poll the future, and ignore the result (either it's `Ok(())`, or it's
-        // `Err(Aborted)`).
-        let _ = future.await;
-    });
-
-    JoinHandle { remote_handle: Some(remote_handle), abort_handle }
-}
+//! On non Wasm platforms, this re-exports parts of tokio directly.  For Wasm,
+//! we provide a single-threaded solution that matches the interface that tokio
+//! provides as a drop in replacement.
 
-#[cfg(target_arch = "wasm32")]
-#[derive(Debug)]
-pub struct JoinHandle<T> {
-    remote_handle: Option<RemoteHandle<T>>,
-    abort_handle: AbortHandle,
+#[cfg(not(target_family = "wasm"))]
+mod sys {
+    pub use tokio::{
+        runtime::{Handle, Runtime},
+        task::{spawn, AbortHandle, JoinError, JoinHandle},
+    };
 }
 
-#[cfg(target_arch = "wasm32")]
-impl<T> JoinHandle<T> {
-    pub fn abort(&self) {
-        self.abort_handle.abort();
+#[cfg(target_family = "wasm")]
+mod sys {
+    use std::{
+        future::Future,
+        pin::Pin,
+        task::{Context, Poll},
+    };
+
+    pub use futures_util::future::AbortHandle;
+    use futures_util::{
+        future::{Abortable, RemoteHandle},
+        FutureExt,
+    };
+
+    /// A Wasm specific version of `tokio::task::JoinError` designed to work
+    /// in the single-threaded environment available in Wasm environments.
+    #[derive(Debug)]
+    pub enum JoinError {
+        Cancelled,
+        Panic,
     }
 
-    pub fn is_finished(&self) -> bool {
-        self.abort_handle.is_aborted()
+    impl JoinError {
+        /// Returns true if the error was caused by the task being cancelled.
+        ///
+        /// See [the module level docs] for more information on cancellation.
+        ///
+        /// [the module level docs]: crate::task#cancellation
+        pub fn is_cancelled(&self) -> bool {
+            matches!(self, JoinError::Cancelled)
+        }
     }
-}
 
-#[cfg(target_arch = "wasm32")]
-impl<T> Drop for JoinHandle<T> {
-    fn drop(&mut self) {
-        // don't abort the spawned future
-        if let Some(h) = self.remote_handle.take() {
-            h.forget();
+    impl std::fmt::Display for JoinError {
+        fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            match &self {
+                JoinError::Cancelled => write!(fmt, "task was cancelled"),
+                JoinError::Panic => write!(fmt, "task panicked"),
+            }
         }
     }
-}
 
-#[cfg(target_arch = "wasm32")]
-impl<T: 'static> Future for JoinHandle<T> {
-    type Output = Result<T, JoinError>;
-
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        if self.abort_handle.is_aborted() {
-            // The future has been aborted. It is not possible to poll it again.
-            Poll::Ready(Err(JoinError))
-        } else if let Some(handle) = self.remote_handle.as_mut() {
-            Pin::new(handle).poll(cx).map(Ok)
-        } else {
-            Poll::Ready(Err(JoinError))
+    /// A Wasm specific version of `tokio::task::JoinHandle` that
+    /// holds handles to locally executing futures.
+    #[derive(Debug)]
+    pub struct JoinHandle<T> {
+        remote_handle: Option<RemoteHandle<T>>,
+        abort_handle: AbortHandle,
+    }
+
+    impl<T> JoinHandle<T> {
+        /// Aborts the spawned future, preventing it from being polled again.
+        pub fn abort(&self) {
+            self.abort_handle.abort();
+        }
+
+        /// Returns the handle to the `AbortHandle` that can be used to
+        /// abort the spawned future.
+        pub fn abort_handle(&self) -> AbortHandle {
+            self.abort_handle.clone()
+        }
+
+        /// Returns true if the spawned future has been aborted.
+        pub fn is_finished(&self) -> bool {
+            self.abort_handle.is_aborted()
         }
     }
+
+    impl<T> Drop for JoinHandle<T> {
+        fn drop(&mut self) {
+            // don't abort the spawned future
+            if let Some(h) = self.remote_handle.take() {
+                h.forget();
+            }
+        }
+    }
+
+    impl<T: 'static> Future for JoinHandle<T> {
+        type Output = Result<T, JoinError>;
+
+        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            if self.abort_handle.is_aborted() {
+                // The future has been aborted. It is not possible to poll it again.
+                Poll::Ready(Err(JoinError::Cancelled))
+            } else if let Some(handle) = self.remote_handle.as_mut() {
+                Pin::new(handle).poll(cx).map(Ok)
+            } else {
+                Poll::Ready(Err(JoinError::Panic))
+            }
+        }
+    }
+
+    /// A Wasm specific version of `tokio::task::spawn` that utilizes
+    /// wasm_bindgen_futures to spawn futures on the local executor.
+    pub fn spawn<F, T>(future: F) -> JoinHandle<T>
+    where
+        F: Future<Output = T> + 'static,
+    {
+        let (future, remote_handle) = future.remote_handle();
+        let (abort_handle, abort_registration) = AbortHandle::new_pair();
+        let future = Abortable::new(future, abort_registration);
+
+        wasm_bindgen_futures::spawn_local(async {
+            // Poll the future, and ignore the result (either it's `Ok(())`, or it's
+            // `Err(Aborted)`).
+            let _ = future.await;
+        });
+
+        JoinHandle { remote_handle: Some(remote_handle), abort_handle }
+    }
 }
 
+pub use sys::*;
+
 #[cfg(test)]
 mod tests {
     use assert_matches::assert_matches;