fix(js): retire dedicated Node workers and bound each worker's dead guest memory

js/bindgen/src/lib.rs

@@ -79,6 +79,8 @@ pub fn set_worker_url(url: String) {
 struct ClientOptions {
     output_bytes: Option<f64>,
     parallelism: Option<f64>,
+    worker_per_thread: Option<bool>,
+    guest_memory_gc_budget: Option<f64>,
     cache: Option<ClientCacheOptions>,
 }
 
@@ -164,7 +166,17 @@ impl JsWasmer {
             .map(|value| validate_usize("parallelism", value, 1))
             .transpose()?
             .and_then(std::num::NonZeroUsize::new);
-        let tasks = Arc::new(tasks::ThreadPool::new(parallelism));
+        let tasks = Arc::new(tasks::ThreadPool::new(
+            parallelism,
+            tasks::PoolOptions {
+                worker_per_thread: options.worker_per_thread.unwrap_or(false),
+                guest_memory_gc_budget: options
+                    .guest_memory_gc_budget
+                    .map(|bytes| validate_usize("guestMemoryGcBudget", bytes, 0))
+                    .transpose()?
+                    .map_or(tasks::DEFAULT_GUEST_MEMORY_GC_BUDGET, |bytes| bytes as u64),
+            },
+        ));
         let mut runtime = PluggableRuntime::new(Arc::clone(&tasks) as Arc<_>);
         runtime
             .set_tty(Arc::new(DefaultTty::default()))

js/bindgen/src/tasks/mod.rs

@@ -48,7 +48,7 @@ pub(crate) use self::{
     post_message_payload::{AsyncJob, BlockingJob, PostMessagePayload},
     scheduler::Scheduler,
     scheduler_message::SchedulerMessage,
-    thread_pool::ThreadPool,
+    thread_pool::{DEFAULT_GUEST_MEMORY_GC_BUDGET, PoolOptions, ThreadPool},
     worker_handle::{CapiTransfer, WorkerHandle},
     worker_message::WorkerMessage,
 };

js/bindgen/src/tasks/post_message_payload.rs

@@ -108,7 +108,10 @@ mod tests {
         let engine = wasmer::Engine::default();
         let module = wasmer::Module::new(&engine, TEST_WASM).unwrap();
         let flag = Arc::new(AtomicBool::new(false));
-        let runtime = PluggableRuntime::new(Arc::new(ThreadPool::new(None)));
+        let runtime = PluggableRuntime::new(Arc::new(ThreadPool::new(
+            None,
+            crate::tasks::PoolOptions::default(),
+        )));
         let env = WasiEnvBuilder::new("program")
             .runtime(Arc::new(runtime))
             .build()

js/bindgen/src/tasks/scheduler.rs

@@ -16,8 +16,8 @@ use wasm_bindgen::JsValue;
 
 use crate::{
     tasks::{
-        AsyncJob, BlockingJob, CapiTransfer, PostMessagePayload, SchedulerMessage, WorkerHandle,
-        WorkerMessage, task_wasm::SpawnWasm, worker_gate::WorkerGate,
+        AsyncJob, BlockingJob, CapiTransfer, PoolOptions, PostMessagePayload, SchedulerMessage,
+        WorkerHandle, WorkerMessage, task_wasm::SpawnWasm, worker_gate::WorkerGate,
     },
     worker_utils::GlobalScope,
 };
@@ -55,14 +55,14 @@ pub(crate) struct Scheduler {
 impl Scheduler {
     /// Spin up a scheduler on the current thread and get a channel that can be
     /// used to communicate with it.
-    pub(crate) fn spawn() -> Scheduler {
+    pub(crate) fn spawn(options: PoolOptions) -> Scheduler {
         let (sender, mut receiver) = mpsc::unbounded_channel();
 
         let thread_id = wasmer::js::current_thread_id();
         // Safety: we just got the thread ID.
         let sender = unsafe { Scheduler::new(sender, thread_id) };
 
-        let mut scheduler = SchedulerState::new(sender.clone());
+        let mut scheduler = SchedulerState::new(sender.clone(), options);
 
         tracing::debug!(thread_id, "Spinning up the scheduler");
         wasm_bindgen_futures::spawn_local(
@@ -225,6 +225,17 @@ struct SchedulerState {
     timer_worker: Option<WorkerHandle>,
     retire_wasm_workers: bool,
     retiring: BTreeSet<u32>,
+    /// Workers retired at their own request because they cannot collect
+    /// their dead guest memory, kept until they report Retired.
+    stopping: BTreeMap<u32, WorkerHandle>,
+    /// Busy workers that asked to retire; they retire instead of going idle.
+    retire_when_idle: BTreeSet<u32>,
+    /// Workers that received a WASIX thread since they were last idle.
+    /// `wasm_workers` cannot tell: every sandbox's main thread has the same
+    /// PID and TID, so an overlapping guest replaces the entry of the one
+    /// still unwinding.
+    ran_wasm: BTreeSet<u32>,
+    guest_memory_gc_budget: u64,
     /// A channel that can be used to send messages to this scheduler.
     mailbox: Scheduler,
     /// Nested WebAssembly objects waiting to travel with the next blocking
@@ -285,19 +296,26 @@ struct DeferredWasm {
 }
 
 impl SchedulerState {
-    fn new(mailbox: Scheduler) -> Self {
+    fn new(mailbox: Scheduler, options: PoolOptions) -> Self {
         SchedulerState {
             idle: VecDeque::new(),
             busy: VecDeque::new(),
             timer_worker: None,
-            retire_wasm_workers: js_sys::Reflect::get(
-                &js_sys::global(),
-                &JsValue::from_str("__wasmerWorkerPerThread"),
-            )
-            .ok()
-            .and_then(|value| value.as_bool())
-            .unwrap_or(false),
+            // The global predates the client option; the Swift WebKit SDK
+            // sets it for every client.
+            retire_wasm_workers: options.worker_per_thread
+                || js_sys::Reflect::get(
+                    &js_sys::global(),
+                    &JsValue::from_str("__wasmerWorkerPerThread"),
+                )
+                .ok()
+                .and_then(|value| value.as_bool())
+                .unwrap_or(false),
             retiring: BTreeSet::new(),
+            stopping: BTreeMap::new(),
+            retire_when_idle: BTreeSet::new(),
+            ran_wasm: BTreeSet::new(),
+            guest_memory_gc_budget: options.guest_memory_gc_budget,
             mailbox,
             pending_capi_transfers: BTreeMap::new(),
             capi_values: BTreeMap::new(),
@@ -313,6 +331,18 @@ impl SchedulerState {
         }
     }
 
+    /// Take an idle worker out of the pool and let it retire once its
+    /// remaining jobs finish (see [`WorkerHandle::retire`]).
+    fn retire_idle_worker(&mut self, worker_id: u32) {
+        let Some(ix) = self.idle.iter().position(|worker| worker.id() == worker_id) else {
+            return;
+        };
+        let worker = self.idle.remove(ix).unwrap();
+        scheduler_diag(format!("retire idle worker={worker_id}"));
+        worker.retire();
+        self.stopping.insert(worker_id, worker);
+    }
+
     /// The worker is gone: free its wasm-bindgen TLS block and stack.
     fn release_worker_memory(&mut self, worker_id: u32) {
         if let Some(gate) = self.gates.remove(&worker_id) {
@@ -337,10 +367,12 @@ impl SchedulerState {
             deferred.task.process().terminate(1.into());
         }
         let mut workers: BTreeMap<u32, WorkerHandle> = BTreeMap::new();
+        let stopping = std::mem::take(&mut self.stopping);
         let all = self
             .idle
             .drain(..)
             .chain(self.busy.drain(..))
+            .chain(stopping.into_values())
             .chain(self.timer_worker.take());
         for worker in all {
             worker.shutdown();
@@ -359,7 +391,12 @@ impl SchedulerState {
                 Either::Right(_) => break,
             };
             match message {
-                Some(SchedulerMessage::WorkerExited { worker_id }) => {
+                // A dedicated worker may retire on its own before it sees
+                // the shutdown request; either way it has stopped.
+                Some(
+                    SchedulerMessage::WorkerExited { worker_id }
+                    | SchedulerMessage::WorkerRetired { worker_id },
+                ) => {
                     if let Some(worker) = workers.remove(&worker_id) {
                         worker.mark_exited();
                     }
@@ -465,9 +502,24 @@ impl SchedulerState {
                 // worker, which was already removed from the pool. Node
                 // reports it once the thread has stopped; browsers do not.
                 tracing::trace!(worker.id = worker_id, "Worker exited");
+                if let Some(worker) = self.stopping.remove(&worker_id) {
+                    worker.mark_exited();
+                }
                 self.release_worker_memory(worker_id);
                 Ok(())
             }
+            SchedulerMessage::WorkerRetireRequested { worker_id } => {
+                if self.idle.iter().any(|worker| worker.id() == worker_id) {
+                    self.retire_idle_worker(worker_id);
+                } else if self.busy.iter().any(|worker| worker.id() == worker_id)
+                    && !self.retiring.contains(&worker_id)
+                {
+                    // It received a task after asking; stop routing work to
+                    // it once that task is done.
+                    self.retire_when_idle.insert(worker_id);
+                }
+                Ok(())
+            }
             SchedulerMessage::WorkerIdle { worker_id } => {
                 // The killed thread unwound on its own; keep the worker.
                 self.doomed.remove(&worker_id);
@@ -482,15 +534,22 @@ impl SchedulerState {
                     .retain(|_, worker| worker.worker_id != worker_id);
                 // Dropping the completed process can expire memories imported
                 // by any worker. Idle workers otherwise retain those JS roots
-                // until they happen to receive another task.
+                // until they happen to receive another task, and this worker
+                // could not release its finished guest's memory.
                 // Timer thunks also emit Idle; they must not broadcast a
                 // collection to the entire pool on every network poll.
-                if self.wasm_workers.len() != active_wasm_threads {
+                let ran_wasm = self.ran_wasm.remove(&worker_id);
+                if ran_wasm || self.wasm_workers.len() != active_wasm_threads {
                     wasmer::js::collect_shared_objects();
                     for worker in &self.idle {
                         worker.collect_shared_objects()?;
                     }
                 }
+                // After its collection message, so the worker drops the
+                // finished guest's roots before it retires.
+                if self.retire_when_idle.remove(&worker_id) {
+                    self.retire_idle_worker(worker_id);
+                }
                 tracing::trace!(
                     worker.id=worker_id,
                     idle_workers=?self.idle.iter().map(|w| w.id()).collect::<Vec<_>>(),
@@ -510,13 +569,19 @@ impl SchedulerState {
             }
             SchedulerMessage::WorkerRetired { worker_id } => {
                 self.doomed.remove(&worker_id);
+                if let Some(worker) = self.stopping.remove(&worker_id) {
+                    worker.mark_exited();
+                    self.release_worker_memory(worker_id);
+                    return Ok(());
+                }
                 if !self.retiring.remove(&worker_id) {
                     if self.worker(worker_id).is_none() {
                         return Ok(());
                     }
                     anyhow::bail!("unexpected retired worker {worker_id}");
                 }
                 self.busy.retain(|worker| worker.id() != worker_id);
+                self.ran_wasm.remove(&worker_id);
                 self.wasm_workers
                     .retain(|_, worker| worker.worker_id != worker_id);
                 self.release_worker_memory(worker_id);
@@ -687,7 +752,11 @@ impl SchedulerState {
             .iter()
             .filter(|worker| {
                 let id = worker.id();
-                if self.retiring.contains(&id) || self.doomed.contains_key(&id) {
+                // A worker that asked to retire goes away instead of idle.
+                if self.retiring.contains(&id)
+                    || self.retire_when_idle.contains(&id)
+                    || self.doomed.contains_key(&id)
+                {
                     return false;
                 }
                 let mut tasks = self
@@ -716,6 +785,7 @@ impl SchedulerState {
             self.next_available_worker(&reason)?
         };
         let transfers = self.take_capi_transfers(source_worker_id, true, Some(worker.id()));
+        self.ran_wasm.insert(worker.id());
         self.wasm_workers.insert(
             task_key,
             WasmWorker {
@@ -769,7 +839,13 @@ impl SchedulerState {
 
         let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
 
-        let handle = WorkerHandle::spawn(id, self.mailbox.clone(), retire_after_task)?;
+        // A dedicated worker's isolate is disposed with its guest memory.
+        let budget = if retire_after_task {
+            0.0
+        } else {
+            self.guest_memory_gc_budget as f64
+        };
+        let handle = WorkerHandle::spawn(id, self.mailbox.clone(), retire_after_task, budget)?;
         self.gates.insert(id, handle.gate());
 
         Ok(handle)
@@ -898,6 +974,8 @@ impl SchedulerState {
     /// Take a worker out of the pool along with everything routed to it.
     fn remove_worker(&mut self, worker_id: u32) -> Option<WorkerHandle> {
         self.retiring.remove(&worker_id);
+        self.retire_when_idle.remove(&worker_id);
+        self.ran_wasm.remove(&worker_id);
         self.wasm_workers
             .retain(|_, worker| worker.worker_id != worker_id);
         self.pending_capi_requests.retain(|_, workers| {
@@ -977,7 +1055,7 @@ mod tests {
         let (sender, receiver) = oneshot::channel();
         let (tx, _) = mpsc::unbounded_channel();
         let tx = unsafe { Scheduler::new(tx, wasmer::js::current_thread_id()) };
-        let mut scheduler = SchedulerState::new(tx);
+        let mut scheduler = SchedulerState::new(tx, PoolOptions::default());
         let message = SchedulerMessage::SpawnAsync(Box::new(move || {
             Box::pin(async move {
                 let _ = sender.send(42);
@@ -1006,7 +1084,7 @@ mod tests {
     fn capi_objects_attach_once_but_remain_available_until_release() {
         let (tx, _) = mpsc::unbounded_channel();
         let tx = unsafe { Scheduler::new(tx, wasmer::js::current_thread_id()) };
-        let mut scheduler = SchedulerState::new(tx);
+        let mut scheduler = SchedulerState::new(tx, PoolOptions::default());
         scheduler
             .pending_capi_transfers
             .entry(7)

js/bindgen/src/tasks/scheduler_message.rs

@@ -45,6 +45,11 @@ pub(crate) enum SchedulerMessage {
     WorkerRetired {
         worker_id: u32,
     },
+    /// An idle worker asks to be retired because it cannot collect the guest
+    /// memory it has released.
+    WorkerRetireRequested {
+        worker_id: u32,
+    },
     /// Terminate the browser worker executing a WASIX thread.
     TerminateWasmThread {
         pid: u32,
@@ -185,6 +190,7 @@ impl SchedulerMessage {
             }
             SchedulerMessage::ForceTerminate { .. }
             | SchedulerMessage::WorkerExited { .. }
+            | SchedulerMessage::WorkerRetireRequested { .. }
             | SchedulerMessage::DrainWaitElapsed { .. } => {
                 Err(anyhow::anyhow!("worker lifecycle messages are local to the scheduler").into())
             }