feat: add channel.onClose (#120)

Author: AriPerkkio

src/common.ts

@@ -3,10 +3,13 @@ import type { MessagePort, TransferListItem } from 'node:worker_threads'
 /** Channel for communicating between main thread and workers */
 export interface TinypoolChannel {
   /** Workers subscribing to messages */
-  onMessage(callback: (message: any) => void): void
+  onMessage?: (callback: (message: any) => void) => void
 
   /** Called with worker's messages */
-  postMessage(message: any): void
+  postMessage?: (message: any) => void
+
+  /** Called when channel can be closed */
+  onClose?: () => void
 }
 
 export interface TinypoolWorker {

src/runtime/process-worker.ts

@@ -64,14 +64,20 @@ export default class ProcessWorker implements TinypoolWorker {
     this.process.stdout?.unpipe(process.stdout)
     this.process.stderr?.unpipe(process.stderr)
     this.port?.close()
+    this.channel?.onClose?.()
     clearTimeout(sigkillTimeout)
   }
 
   setChannel(channel: TinypoolChannel) {
+    // Previous channel exists in non-isolated runs
+    if (this.channel && this.channel !== channel) {
+      this.channel.onClose?.()
+    }
+
     this.channel = channel
 
     // Mirror channel's messages to process
-    this.channel.onMessage((message: any) => {
+    this.channel.onMessage?.((message: any) => {
       this.send(message)
     })
   }
@@ -115,7 +121,7 @@ export default class ProcessWorker implements TinypoolWorker {
       }
 
       if (!data || !data.__tinypool_worker_message__) {
-        return this.channel?.postMessage(data)
+        return this.channel?.postMessage?.(data)
       }
 
       if (data.source === 'pool') {

src/runtime/thread-worker.ts

@@ -1,12 +1,13 @@
 import { fileURLToPath } from 'node:url'
 import { type TransferListItem, Worker } from 'node:worker_threads'
-import { type TinypoolWorker } from '../common'
+import { type TinypoolWorker, type TinypoolChannel } from '../common'
 
 export default class ThreadWorker implements TinypoolWorker {
   name = 'ThreadWorker'
   runtime = 'worker_threads'
   thread!: Worker
   threadId!: number
+  channel?: TinypoolChannel
 
   initialize(options: Parameters<TinypoolWorker['initialize']>[0]) {
     this.thread = new Worker(
@@ -17,7 +18,11 @@ export default class ThreadWorker implements TinypoolWorker {
   }
 
   async terminate() {
-    return this.thread.terminate()
+    const output = await this.thread.terminate()
+
+    this.channel?.onClose?.()
+
+    return output
   }
 
   postMessage(message: any, transferListItem?: Readonly<TransferListItem[]>) {
@@ -44,9 +49,24 @@ export default class ThreadWorker implements TinypoolWorker {
     return this.thread.unref()
   }
 
-  setChannel() {
-    throw new Error(
-      "{ runtime: 'worker_threads' } doesn't support channel. Use transferListItem instead."
-    )
+  setChannel(channel: TinypoolChannel) {
+    if (channel.onMessage) {
+      throw new Error(
+        "{ runtime: 'worker_threads' } doesn't support channel.onMessage. Use transferListItem for listening to messages instead."
+      )
+    }
+
+    if (channel.postMessage) {
+      throw new Error(
+        "{ runtime: 'worker_threads' } doesn't support channel.postMessage. Use transferListItem for sending to messages instead."
+      )
+    }
+
+    // Previous channel exists in non-isolated runs
+    if (this.channel && this.channel !== channel) {
+      this.channel.onClose?.()
+    }
+
+    this.channel = channel
   }
 }

test/runtime.test.ts

@@ -46,6 +46,52 @@ describe('worker_threads', () => {
     `)
     expect(result).toBe(threadId)
   })
+
+  test('channel is closed when isolated', async () => {
+    const pool = createPool({
+      runtime: 'worker_threads',
+      isolateWorkers: true,
+      minThreads: 2,
+      maxThreads: 2,
+    })
+
+    const events: string[] = []
+
+    await pool.run('', { channel: { onClose: () => events.push('call #1') } })
+    expect(events).toStrictEqual(['call #1'])
+
+    await pool.run('', { channel: { onClose: () => events.push('call #2') } })
+    expect(events).toStrictEqual(['call #1', 'call #2'])
+
+    await pool.run('', { channel: { onClose: () => events.push('call #3') } })
+    expect(events).toStrictEqual(['call #1', 'call #2', 'call #3'])
+
+    await pool.destroy()
+    expect(events).toStrictEqual(['call #1', 'call #2', 'call #3'])
+  })
+
+  test('channel is closed when non-isolated', async () => {
+    const pool = createPool({
+      runtime: 'worker_threads',
+      isolateWorkers: false,
+      minThreads: 2,
+      maxThreads: 2,
+    })
+
+    const events: string[] = []
+
+    await pool.run('', { channel: { onClose: () => events.push('call #1') } })
+    expect(events).toStrictEqual([])
+
+    await pool.run('', { channel: { onClose: () => events.push('call #2') } })
+    expect(events).toStrictEqual(['call #1'])
+
+    await pool.run('', { channel: { onClose: () => events.push('call #3') } })
+    expect(events).toStrictEqual(['call #1', 'call #2'])
+
+    await pool.destroy()
+    expect(events).toStrictEqual(['call #1', 'call #2', 'call #3'])
+  })
 })
 
 describe('child_process', () => {
@@ -152,6 +198,52 @@ describe('child_process', () => {
       response: 'Hello from worker',
     })
   })
+
+  test('channel is closed when isolated', async () => {
+    const pool = createPool({
+      runtime: 'child_process',
+      isolateWorkers: true,
+      minThreads: 2,
+      maxThreads: 2,
+    })
+
+    const events: string[] = []
+
+    await pool.run('', { channel: { onClose: () => events.push('call #1') } })
+    expect(events).toStrictEqual(['call #1'])
+
+    await pool.run('', { channel: { onClose: () => events.push('call #2') } })
+    expect(events).toStrictEqual(['call #1', 'call #2'])
+
+    await pool.run('', { channel: { onClose: () => events.push('call #3') } })
+    expect(events).toStrictEqual(['call #1', 'call #2', 'call #3'])
+
+    await pool.destroy()
+    expect(events).toStrictEqual(['call #1', 'call #2', 'call #3'])
+  })
+
+  test('channel is closed when non-isolated', async () => {
+    const pool = createPool({
+      runtime: 'child_process',
+      isolateWorkers: false,
+      minThreads: 2,
+      maxThreads: 2,
+    })
+
+    const events: string[] = []
+
+    await pool.run('', { channel: { onClose: () => events.push('call #1') } })
+    expect(events).toStrictEqual([])
+
+    await pool.run('', { channel: { onClose: () => events.push('call #2') } })
+    expect(events).toStrictEqual(['call #1'])
+
+    await pool.run('', { channel: { onClose: () => events.push('call #3') } })
+    expect(events).toStrictEqual(['call #1', 'call #2'])
+
+    await pool.destroy()
+    expect(events).toStrictEqual(['call #1', 'call #2', 'call #3'])
+  })
 })
 
 test('runtime can be changed after recycle', async () => {