From 0f83592121efac86967c15c1bcf2829ca994c843 Mon Sep 17 00:00:00 2001
From: AriPerkkio <ari.perkkio@gmail.com>
Date: Fri, 12 Apr 2024 06:17:20 +0300
Subject: [PATCH] fix(child_process): prevent writing to terminating process

---
 src/runtime/process-worker.ts | 14 +++++++++++---
 test/termination.test.ts      | 26 ++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/runtime/process-worker.ts b/src/runtime/process-worker.ts
index fddd214..4c07919 100644
--- a/src/runtime/process-worker.ts
+++ b/src/runtime/process-worker.ts
@@ -18,6 +18,7 @@ export default class ProcessWorker implements TinypoolWorker {
   port?: MessagePort
   channel?: TinypoolChannel
   waitForExit!: Promise<void>
+  isTerminating = false
 
   initialize(options: Parameters<TinypoolWorker['initialize']>[0]) {
     this.process = fork(
@@ -42,6 +43,7 @@ export default class ProcessWorker implements TinypoolWorker {
   }
 
   async terminate() {
+    this.isTerminating = true
     this.process.off('exit', this.onUnexpectedExit)
 
     const sigkillTimeout = setTimeout(
@@ -61,10 +63,16 @@ export default class ProcessWorker implements TinypoolWorker {
 
     // Mirror channel's messages to process
     this.channel.onMessage((message: any) => {
-      this.process.send(message)
+      this.send(message)
     })
   }
 
+  private send(message: Parameters<NonNullable<typeof process['send']>>[0]) {
+    if (!this.isTerminating) {
+      this.process.send(message)
+    }
+  }
+
   postMessage(message: any, transferListItem?: Readonly<TransferListItem[]>) {
     transferListItem?.forEach((item) => {
       if (item instanceof MessagePort) {
@@ -75,7 +83,7 @@ export default class ProcessWorker implements TinypoolWorker {
     // Mirror port's messages to process
     if (this.port) {
       this.port.on('message', (message) =>
-        this.process.send(<TinypoolWorkerMessage<'port'>>{
+        this.send(<TinypoolWorkerMessage<'port'>>{
           ...message,
           source: 'port',
           __tinypool_worker_message__,
@@ -83,7 +91,7 @@ export default class ProcessWorker implements TinypoolWorker {
       )
     }
 
-    return this.process.send(<TinypoolWorkerMessage<'pool'>>{
+    return this.send(<TinypoolWorkerMessage<'pool'>>{
       ...message,
       source: 'pool',
       __tinypool_worker_message__,
diff --git a/test/termination.test.ts b/test/termination.test.ts
index e6146ea..b34f922 100644
--- a/test/termination.test.ts
+++ b/test/termination.test.ts
@@ -30,3 +30,29 @@ test('termination timeout throws when worker does not terminate in time', async
     'Failed to terminate worker'
   )
 })
+
+test('writing to terminating worker does not crash', async () => {
+  const listeners: ((msg: any) => void)[] = []
+
+  const pool = new Tinypool({
+    runtime: 'child_process',
+    filename: resolve(__dirname, 'fixtures/sleep.js'),
+    minThreads: 1,
+    maxThreads: 1,
+  })
+
+  await pool.run(
+    {},
+    {
+      channel: {
+        onMessage: (listener) => listeners.push(listener),
+        postMessage: () => {},
+      },
+    }
+  )
+
+  const destroyed = pool.destroy()
+  listeners.forEach((listener) => listener('Hello from main thread'))
+
+  await destroyed
+})