feat: add teardown option (#115)

Author: AriPerkkio

README.md

@@ -180,6 +180,7 @@ We have a similar API to Piscina, so for more information, you can read Piscina'
 - `runtime`: Used to pick worker runtime. Default value is `worker_threads`.
   - `worker_threads`: Runs workers in [`node:worker_threads`](https://nodejs.org/api/worker_threads.html). For `main thread <-> worker thread` communication you can use [`MessagePort`](https://nodejs.org/api/worker_threads.html#class-messageport) in the `pool.run()` method's [`transferList` option](https://nodejs.org/api/worker_threads.html#portpostmessagevalue-transferlist). See [example](#main-thread---worker-thread-communication).
   - `child_process`: Runs workers in [`node:child_process`](https://nodejs.org/api/child_process.html). For `main thread <-> worker process` communication you can use `TinypoolChannel` in the `pool.run()` method's `channel` option. For filtering out the Tinypool's internal messages see `TinypoolWorkerMessage`. See [example](#main-process---worker-process-communication).
+- `teardown`: name of the function in file that should be called before worker is terminated. Must be named exported.
 
 #### Pool methods
 

package.json

@@ -37,7 +37,7 @@
   "scripts": {
     "test": "vitest",
     "bench": "vitest bench",
-    "dev": "tsdown --watch",
+    "dev": "tsdown --watch ./src",
     "build": "tsdown",
     "publish": "clean-publish",
     "lint": "eslint --max-warnings=0",

src/index.ts

@@ -152,6 +152,7 @@ interface Options {
   taskQueue?: TaskQueue
   trackUnmanagedFds?: boolean
   isolateWorkers?: boolean
+  teardown?: string
 }
 
 interface FilledOptions extends Options {
@@ -454,18 +455,24 @@ class WorkerInfo extends AsynchronouslyCreatedResource {
   usedMemory?: number
   onMessage: ResponseCallback
   shouldRecycle?: boolean
+  filename?: string | null
+  teardown?: string
 
   constructor(
     worker: TinypoolWorker,
     port: MessagePort,
     workerId: number,
     freeWorkerId: () => void,
-    onMessage: ResponseCallback
+    onMessage: ResponseCallback,
+    filename?: string | null,
+    teardown?: string
   ) {
     super()
     this.worker = worker
     this.workerId = workerId
     this.freeWorkerId = freeWorkerId
+    this.teardown = teardown
+    this.filename = filename
     this.port = port
     this.port.on('message', (message: ResponseMessage) =>
       this._handleResponse(message)
@@ -486,6 +493,25 @@ class WorkerInfo extends AsynchronouslyCreatedResource {
       reject = rej
     })
 
+    if (this.teardown && this.filename) {
+      const { teardown, filename } = this
+
+      await new Promise((resolve, reject) => {
+        this.postTask(
+          new TaskInfo(
+            {},
+            [],
+            filename,
+            teardown,
+            (error, result) => (error ? reject(error) : resolve(result)),
+            null,
+            1,
+            undefined
+          )
+        )
+      })
+    }
+
     const timer = timeout
       ? setTimeout(
           () => reject(new Error('Failed to terminate worker')),
@@ -743,7 +769,9 @@ class ThreadPool {
       port1,
       workerId!,
       () => workerIds.set(workerId, true),
-      onMessage
+      onMessage,
+      this.options.filename,
+      this.options.teardown
     )
     if (this.startingUp) {
       // There is no point in waiting for the initial set of Workers to indicate

test/fixtures/teardown.mjs

@@ -0,0 +1,19 @@
+import { setTimeout } from 'node:timers/promises'
+
+let state = 0
+
+/** @type {import("node:worker_threads").MessagePort } */
+let port
+
+export default function task(options) {
+  port ||= options?.port
+  state++
+
+  return `Output of task #${state}`
+}
+
+export async function namedTeardown() {
+  await setTimeout(50)
+
+  port?.postMessage(`Teardown of task #${state}`)
+}

test/teardown.test.ts

@@ -0,0 +1,59 @@
+import { dirname, resolve } from 'node:path'
+import { Tinypool } from 'tinypool'
+import { fileURLToPath } from 'node:url'
+import { MessageChannel } from 'node:worker_threads'
+
+const __dirname = dirname(fileURLToPath(import.meta.url))
+
+test('isolated workers call teardown on worker recycle', async () => {
+  const pool = new Tinypool({
+    filename: resolve(__dirname, 'fixtures/teardown.mjs'),
+    minThreads: 1,
+    maxThreads: 1,
+    isolateWorkers: true,
+    teardown: 'namedTeardown',
+  })
+
+  for (const _ of [1, 2, 3, 4, 5]) {
+    const { port1, port2 } = new MessageChannel()
+    const promise = new Promise((resolve) => port2.on('message', resolve))
+
+    const output = await pool.run({ port: port1 }, { transferList: [port1] })
+    expect(output).toBe('Output of task #1')
+
+    await expect(promise).resolves.toBe('Teardown of task #1')
+  }
+})
+
+test('non-isolated workers call teardown on worker recycle', async () => {
+  const pool = new Tinypool({
+    filename: resolve(__dirname, 'fixtures/teardown.mjs'),
+    minThreads: 1,
+    maxThreads: 1,
+    isolateWorkers: false,
+    teardown: 'namedTeardown',
+  })
+
+  function unexpectedTeardown(message: string) {
+    assert.fail(
+      `Teardown should not have been called yet. Received "${message}"`
+    )
+  }
+
+  const { port1, port2 } = new MessageChannel()
+
+  for (const index of [1, 2, 3, 4, 5]) {
+    port2.on('message', unexpectedTeardown)
+
+    const transferList = index === 1 ? [port1] : []
+
+    const output = await pool.run({ port: transferList[0] }, { transferList })
+    expect(output).toBe(`Output of task #${index}`)
+  }
+
+  port2.off('message', unexpectedTeardown)
+  const promise = new Promise((resolve) => port2.on('message', resolve))
+
+  await pool.destroy()
+  await expect(promise).resolves.toMatchInlineSnapshot(`"Teardown of task #5"`)
+})