fix: cleanup worker threads channel when project shutdown

Author: xusd320

crates/napi/src/next_api/project.rs

@@ -722,6 +722,7 @@ pub async fn project_shutdown(
 ) {
     project.turbopack_ctx.turbo_tasks().stop_and_wait().await;
     project_on_exit_internal(&project).await;
+    turbopack_node::worker_pool::shutdown();
 }
 
 #[napi(object)]

packages/next/src/build/swc/index.ts

@@ -671,38 +671,54 @@ function bindingToApi(
   const loaderWorkers: Record<string, Array<Worker>> = {}
 
   const createOrScalePool = async () => {
-    let poolOptions = await binding.recvPoolRequest()
-    const { filename, maxConcurrency } = poolOptions
-    const workers = loaderWorkers[filename] || (loaderWorkers[filename] = [])
-    if (workers.length < maxConcurrency) {
-      for (let i = workers.length; i < maxConcurrency; i++) {
-        const worker = new Worker(filename, {
-          workerData: {
-            poolId: filename,
-            bindingPath,
-          },
-        })
-        worker.unref()
-        workers.push(worker)
+    while (true) {
+      try {
+        let poolOptions = await binding.recvPoolRequest()
+        const { filename, maxConcurrency } = poolOptions
+        const workers =
+          loaderWorkers[filename] || (loaderWorkers[filename] = [])
+        if (workers.length < maxConcurrency) {
+          for (let i = workers.length; i < maxConcurrency; i++) {
+            const worker = new Worker(filename, {
+              workerData: {
+                poolId: filename,
+                bindingPath,
+              },
+            })
+            worker.unref()
+            workers.push(worker)
+          }
+        } else if (workers.length > maxConcurrency) {
+          const workersToStop = workers.splice(
+            0,
+            workers.length - maxConcurrency
+          )
+          workersToStop.forEach((worker) => worker.terminate())
+        }
+      } catch (_) {
+        // rust channel closed, do nothing
+        return
       }
-    } else if (workers.length > maxConcurrency) {
-      const workersToStop = workers.splice(0, workers.length - maxConcurrency)
-      workersToStop.forEach((worker) => worker.terminate())
     }
-    createOrScalePool()
   }
 
   const waitingForWorkerTermination = async () => {
-    const { filename, workerId } = await binding.recvWorkerTermination()
-    const workers = loaderWorkers[filename]
-    const workerIdx = workers.findIndex(
-      (worker) => worker.threadId === workerId
-    )
-    if (workerIdx > -1) {
-      const worker = workers.splice(workerIdx, 1)
-      worker[0].terminate()
+    while (true) {
+      try {
+        const { filename, workerId } = await binding.recvWorkerTermination()
+        const workers = loaderWorkers[filename]
+        const workerIdx = workers.findIndex(
+          (worker) => worker.threadId === workerId
+        )
+        if (workerIdx > -1) {
+          const worker = workers.splice(workerIdx, 1)
+          worker[0].terminate()
+        }
+      } catch (_) {
+        // rust channel closed, do nothing
+        return
+      }
     }
-    waitingForWorkerTermination()
   }
 
   class ProjectImpl implements Project {

turbopack/crates/turbopack-node/js/src/worker_threads/evaluate.ts

@@ -121,6 +121,9 @@ export const run = async (
   }
 
   while (true) {
+    // need await a micro task, or else if some request rejected,
+    // the error will be propergated to schedule thread, then causing panic
+    await Promise.resolve()
     const msg_str = await binding.recvMessageInWorker(workerId)
 
     const msg = JSON.parse(msg_str) as
@@ -150,9 +153,6 @@ export const run = async (
           requests.delete(msg.id)
           if (msg.error) {
             request.reject(new Error(msg.error))
-            // need await a micro task, or else if some request rejected,
-            // the error will be propergated to schedule thread, then causing panic
-            await Promise.resolve()
           } else {
             request.resolve(msg.data)
           }

turbopack/crates/turbopack-node/src/lib.rs

@@ -23,7 +23,7 @@ mod process_pool;
 pub mod source_map;
 pub mod transforms;
 #[cfg(feature = "worker_pool")]
-mod worker_pool;
+pub mod worker_pool;
 
 #[turbo_tasks::function]
 async fn emit(

turbopack/crates/turbopack-node/src/worker_pool/mod.rs

@@ -16,23 +16,24 @@ use crate::{
 };
 
 mod operation;
+pub use operation::shutdown;
 mod worker_thread;
 
 static OPERATION_TASK_ID: AtomicU32 = AtomicU32::new(1);
 
 #[turbo_tasks::value]
-pub struct WorkerThreadPool {
+pub(crate) struct WorkerThreadPool {
     cwd: PathBuf,
     entrypoint: PathBuf,
     env: FxHashMap<RcStr, RcStr>,
     concurrency: usize,
-    pub assets_for_source_mapping: ResolvedVc<AssetsForSourceMapping>,
-    pub assets_root: FileSystemPath,
-    pub project_dir: FileSystemPath,
+    pub(crate) assets_for_source_mapping: ResolvedVc<AssetsForSourceMapping>,
+    pub(crate) assets_root: FileSystemPath,
+    pub(crate) project_dir: FileSystemPath,
 }
 
 impl WorkerThreadPool {
-    pub fn create(
+    pub(crate) fn create(
         cwd: PathBuf,
         entrypoint: PathBuf,
         env: FxHashMap<RcStr, RcStr>,

turbopack/crates/turbopack-node/src/worker_pool/operation.rs

@@ -27,6 +27,11 @@ impl<T: Send + Sync + 'static> MessageChannel<T> {
     pub(crate) async fn recv(&self) -> Result<T> {
         Ok(self.receiver.recv().await?)
     }
+
+    pub(crate) fn close(&self) {
+        self.sender.close();
+        self.receiver.close();
+    }
 }
 
 pub(crate) struct WorkerPoolOperation {
@@ -138,6 +143,13 @@ impl WorkerPoolOperation {
             .context("failed to recv pool request")
     }
 
+    pub(crate) fn shutdown(&self) {
+        // We need to close channels connected to schedule thread,
+        // or else, it will be forever waiting in schedule thread
+        self.pool_request_channel.close();
+        self.worker_termination_channel.close();
+    }
+
     pub(crate) async fn recv_worker_request(&self, pool_id: String) -> Result<u32> {
         let channel = self
             .worker_request_channel
@@ -217,6 +229,10 @@ pub async fn recv_task_message(task_id: u32) -> Result<String> {
     WORKER_POOL_OPERATION.recv_task_response(task_id).await
 }
 
+pub fn shutdown() {
+    WORKER_POOL_OPERATION.shutdown();
+}
+
 pub(crate) struct WorkerOperation {
     pub(crate) pool_id: String,
     pub(crate) task_id: u32,