Report worker termination to channel tasks

README.md

@@ -201,9 +201,36 @@ let progress: Progress = task.recv().await.unwrap();
 task.send(&Continue { should_continue: true });
 
 // Wait for task completion
-let result = task.result().await;
+let result = task.result().await?;
 ```
 
+#### Stopping a task
+
+A channel task is best stopped cooperatively, by sending it a message that its
+function checks for — the `Continue { should_continue: false }` message above is
+an example. The worker then unwinds normally and stays available for the next task.
+
+Work that does not cooperate — a computation without a cancellation point, or a
+task that spawned background work of its own — can only be stopped by terminating
+its worker. Run such a task on its own `WebWorker` and terminate that:
+
+```rust,ignore
+let worker = WebWorker::new(None).await?;
+let task = worker
+   .run_channel(webworker_channel!(process_with_progress), &data)
+   .await;
+
+worker.terminate();
+
+// recv() now returns None and result() reports the termination.
+assert_eq!(task.result().await, Err(TaskError::WorkerTerminated));
+```
+
+Terminating discards the worker, so this is a last resort rather than a routine
+cancellation: the next task pays for the creation of a new worker. Plain `run`
+tasks that are still in flight on that worker cannot report an error and panic
+instead, so only terminate a worker once its plain tasks have completed.
+
 ### Bundler support (Vite)
 The recommended approach for Vite is to place the wasm-pack output in Vite's `publicDir`.
 This keeps the glue code and WASM binary as static assets, which is required because each

src/channel_task.rs

@@ -1,9 +1,13 @@
 use std::marker::PhantomData;
 
+use futures::{
+    future::{select, Either},
+    pin_mut,
+};
 use serde::{de::DeserializeOwned, Serialize};
-use tokio::sync::oneshot;
+use tokio::sync::{oneshot, watch};
 
-use crate::{channel::Channel, convert::from_bytes};
+use crate::{channel::Channel, convert::from_bytes, error::TaskError};
 
 /// A handle to a running channel task on a WebWorker.
 ///
@@ -14,6 +18,9 @@ use crate::{channel::Channel, convert::from_bytes};
 /// [`crate::pool::WebWorkerPool::run_channel`]. It allows you to exchange messages
 /// with the worker (e.g., for progress reporting) and then consume the final result.
 ///
+/// If the worker is terminated while the task is still running, [`ChannelTask::recv`]
+/// returns `None` and [`ChannelTask::result`] returns [`TaskError::WorkerTerminated`].
+///
 /// # Example
 ///
 /// ```ignore
@@ -24,38 +31,57 @@ use crate::{channel::Channel, convert::from_bytes};
 /// let progress: Progress = task.recv().await.expect("progress");
 /// task.send(&Continue { should_continue: true });
 ///
-/// let result: ProcessResult = task.result().await;
+/// let result: ProcessResult = task.result().await.expect("worker terminated");
 /// ```
 pub struct ChannelTask<R> {
     channel: Channel,
     result_rx: oneshot::Receiver<Vec<u8>>,
+    /// Set to `true` once the worker running this task has been terminated.
+    terminated: watch::Receiver<bool>,
     _phantom: PhantomData<R>,
 }
 
 impl<R: DeserializeOwned> ChannelTask<R> {
-    /// Create a new `ChannelTask` from a channel and a result receiver.
-    #[doc(hidden)]
-    pub fn new(channel: Channel, result_rx: oneshot::Receiver<Vec<u8>>) -> Self {
+    /// Create a new `ChannelTask` from a channel, a result receiver, and the
+    /// termination signal of the worker running the task.
+    pub(crate) fn new(
+        channel: Channel,
+        result_rx: oneshot::Receiver<Vec<u8>>,
+        terminated: watch::Receiver<bool>,
+    ) -> Self {
         Self {
             channel,
             result_rx,
+            terminated,
             _phantom: PhantomData,
         }
     }
 
     /// Receive the next deserialized message from the worker.
     ///
-    /// Returns `None` if the channel's sender side has been dropped
-    /// (i.e., the worker has finished and closed the channel).
+    /// Returns `None` once the worker has been terminated and all messages it
+    /// already sent have been received.
     pub async fn recv<T: DeserializeOwned>(&self) -> Option<T> {
-        self.channel.recv().await
+        let bytes = self.recv_bytes().await?;
+        Some(from_bytes(&bytes))
     }
 
     /// Receive raw bytes from the worker.
     ///
-    /// Returns `None` if the channel's sender side has been dropped.
+    /// Returns `None` once the worker has been terminated and all messages it
+    /// already sent have been received.
     pub async fn recv_bytes(&self) -> Option<Box<[u8]>> {
-        self.channel.recv_bytes().await
+        // Messages that already arrived are handed out before reporting the
+        // termination, so no message is lost when a worker is terminated.
+        let mut terminated = self.terminated.clone();
+        let message = self.channel.recv_bytes();
+        let terminated = terminated.changed();
+        pin_mut!(message, terminated);
+
+        match select(message, terminated).await {
+            Either::Left((message, _)) => message,
+            Either::Right(_) => None,
+        }
     }
 
     /// Send a serialized message to the worker.
@@ -69,11 +95,14 @@ impl<R: DeserializeOwned> ChannelTask<R> {
     }
 
     /// Await the task's final result, consuming the `ChannelTask`.
-    pub async fn result(self) -> R {
+    ///
+    /// Returns [`TaskError::WorkerTerminated`] if the worker was terminated
+    /// before the task returned a result.
+    pub async fn result(self) -> Result<R, TaskError> {
         let bytes = self
             .result_rx
             .await
-            .expect("WebWorker result sender dropped");
-        from_bytes(&bytes)
+            .map_err(|_| TaskError::WorkerTerminated)?;
+        Ok(from_bytes(&bytes))
     }
 }

src/error.rs

@@ -8,6 +8,15 @@ use thiserror::Error;
 #[error("WebWorker capacity reached")]
 pub struct Full;
 
+/// This error is returned when a channel task cannot produce a result.
+#[derive(Debug, Error, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum TaskError {
+    /// The worker running the task was terminated before it returned a result.
+    #[error("WebWorker was terminated")]
+    WorkerTerminated,
+}
+
 /// This error is returned during the creation of a new web worker.
 /// It covers generic errors in the actual creation and import errors
 /// during the initialization.

src/lib.rs

@@ -97,6 +97,7 @@
 #![allow(clippy::borrowed_box)]
 pub use channel::Channel;
 pub use channel_task::ChannelTask;
+pub use error::TaskError;
 pub use global::{
     has_worker_pool, init_optimized_worker_pool, init_worker_pool, worker_pool, AlreadyInitialized,
 };

src/pool/mod.rs

@@ -1,3 +1,7 @@
+// The getters that `#[wasm_bindgen(getter_with_clone)]` generates for
+// `WorkerPoolOptions` clone every field, including the `Copy` ones.
+#![allow(clippy::clone_on_copy)]
+
 use std::{borrow::Borrow, cell::RefCell, rc::Rc};
 
 use futures::future::join_all;
@@ -306,7 +310,7 @@ impl WebWorkerPool {
     ///
     /// let progress: Progress = task.recv().await.expect("progress");
     /// task.send(&Continue { should_continue: true });
-    /// let result: ProcessResult = task.result().await;
+    /// let result: ProcessResult = task.result().await.expect("worker terminated");
     /// ```
     pub async fn run_channel<T, R>(&self, func: WebWorkerChannelFn<T, R>, arg: &T) -> ChannelTask<R>
     where

src/webworker/worker.rs

@@ -9,8 +9,8 @@ use super::com::*;
 use super::js::*;
 use js_sys::Array;
 use serde::{Deserialize, Serialize};
-use tokio::sync::{oneshot, Semaphore};
-use wasm_bindgen::{prelude::Closure, JsCast, JsValue, UnwrapThrowExt};
+use tokio::sync::{oneshot, watch, Semaphore};
+use wasm_bindgen::{prelude::Closure, throw_str, JsCast, JsValue, UnwrapThrowExt};
 use web_sys::{
     Blob, BlobPropertyBag, MessageChannel, MessageEvent, MessagePort, Url, Worker, WorkerOptions,
     WorkerType,
@@ -68,6 +68,9 @@ pub struct WebWorker {
     _callback: Closure<Callback>,
     /// Timestamp (ms since epoch) of the last completed task, used for idle timeout tracking.
     last_active: Rc<Cell<f64>>,
+    /// Set to `true` once this worker has been terminated. Channel tasks subscribe
+    /// to this signal so that they can report the termination to their caller.
+    terminated: watch::Sender<bool>,
 }
 
 impl WebWorker {
@@ -220,6 +223,7 @@ impl WebWorker {
             open_tasks: tasks,
             _callback: callback_handle,
             last_active,
+            terminated: watch::channel(false).0,
         })
     }
 
@@ -287,7 +291,7 @@ impl WebWorker {
     ///
     /// let progress: Progress = task.recv().await.expect("progress");
     /// task.send(&Continue { should_continue: true });
-    /// let result: ProcessResult = task.result().await;
+    /// let result: ProcessResult = task.result().await.expect("worker terminated");
     /// ```
     pub async fn run_channel<T, R>(&self, func: WebWorkerChannelFn<T, R>, arg: &T) -> ChannelTask<R>
     where
@@ -430,7 +434,7 @@ impl WebWorker {
         // Send the request and get a receiver for the result bytes.
         let result_rx = self.send_channel_request(func.name, arg, worker_port);
 
-        ChannelTask::new(channel, result_rx)
+        ChannelTask::new(channel, result_rx, self.terminated.subscribe())
     }
 
     /// This function handles the communication with the worker
@@ -462,6 +466,10 @@ impl WebWorker {
     /// Sends a request to the worker and waits for the response.
     /// This is extracted from `force_run` to reduce monomorphisation cost.
     async fn send_request(&self, id: u32, request: Request, port: Option<MessagePort>) -> Vec<u8> {
+        if self.is_terminated() {
+            throw_str("WebWorker has been terminated");
+        }
+
         // Create channel and add task.
         let (sender, receiver) = oneshot::channel();
         self.open_tasks.borrow_mut().insert(id, sender);
@@ -490,7 +498,7 @@ impl WebWorker {
         // Handle result.
         receiver
             .await
-            .expect_throw("WebWorker gone")
+            .expect_throw("WebWorker was terminated before the task completed")
             .response
             .expect_throw("Could not find function")
     }
@@ -507,6 +515,10 @@ impl WebWorker {
     where
         T: Serialize + for<'de> Deserialize<'de>,
     {
+        if self.is_terminated() {
+            throw_str("WebWorker has been terminated");
+        }
+
         let id = self.current_task.fetch_add(1, Ordering::Relaxed);
         let request = Request {
             id,
@@ -554,11 +566,45 @@ impl WebWorker {
     pub fn last_active(&self) -> f64 {
         self.last_active.get()
     }
+
+    /// Terminate this worker.
+    ///
+    /// This is the only way to stop work that does not cooperate, for example a
+    /// long-running computation without a cancellation point, or a task that
+    /// spawned background work of its own. Prefer cooperative cancellation over
+    /// a [`crate::Channel`] where the task supports it: terminating discards the
+    /// worker, so the next task pays for the creation of a new one.
+    ///
+    /// Tasks that are still running on this worker are abandoned: pending
+    /// [`ChannelTask::result`] calls resolve to
+    /// [`crate::TaskError::WorkerTerminated`] and pending [`ChannelTask::recv`]
+    /// calls return `None`. Pending [`WebWorker::run`] calls cannot report an
+    /// error and panic instead, so terminate a worker only once the plain tasks
+    /// on it have completed.
+    ///
+    /// The worker cannot be used afterwards; further calls to `run` or
+    /// `run_channel` panic. Repeated calls to `terminate` are harmless.
+    /// Dropping a [`WebWorker`] terminates it as well.
+    pub fn terminate(&self) {
+        if self.terminated.send_replace(true) {
+            return;
+        }
+
+        self.port.close();
+        self.worker.terminate();
+
+        // Fail all tasks that were still running on this worker.
+        self.open_tasks.borrow_mut().clear();
+    }
+
+    /// Return whether this worker has been terminated.
+    pub fn is_terminated(&self) -> bool {
+        *self.terminated.borrow()
+    }
 }
 
 impl Drop for WebWorker {
     fn drop(&mut self) {
-        self.port.close();
-        self.worker.terminate();
+        self.terminate();
     }
 }