feat: pool.cancelPendingTasks method (#53)

README.md

@@ -42,13 +42,21 @@ export default ({ a, b }) => {
 
 We have a similar API to Piscina, so for more information, you can read Piscina's detailed [documentation](https://github.com/piscinajs/piscina#piscina---the-nodejs-worker-pool) and apply the same techniques here.
 
-###### Additional Options
+### Tinypool specific APIs
+
+#### Pool constructor options
 
 - `isolateWorkers`: Default to `false`. Always starts with a fresh worker when running tasks to isolate the environment.
-- `workerId`: Each worker now has an id ( <= `maxThreads`) that can be imported
-  from `tinypool` in the worker itself (or `process.__tinypool_state__.workerId`)
 - `terminateTimeout`: Defaults to `null`. If terminating a worker takes `terminateTimeout` amount of milliseconds to execute, an error is raised.
 
+#### Pool methods
+
+- `cancelPendingTasks()`: Gracefully cancels all pending tasks without stopping or interfering with on-going tasks. This method is useful when your tasks may have side effects and should not be terminated forcefully during task execution. If your tasks don't have any side effects you may want to use [`{ signal }`](https://github.com/piscinajs/piscina#cancelable-tasks) option for forcefully terminating all tasks, including the on-going ones, instead.
+
+#### Exports
+
+- `workerId`: Each worker now has an id ( <= `maxThreads`) that can be imported from `tinypool` in the worker itself (or `process.__tinypool_state__.workerId`).
+
 ## Authors
 
 | <a href="https://github.com/Aslemammad"> <img width='150' src="https://avatars.githubusercontent.com/u/37929992?v=4" /><br> Mohammad Bagher </a> |

src/common.ts

@@ -70,13 +70,15 @@ export interface Transferable {
 
 export interface Task {
   readonly [kQueueOptions]: object | null
+  cancel(): void
 }
 
 export interface TaskQueue {
   readonly size: number
   shift(): Task | null
   remove(task: Task): void
   push(task: Task): void
+  cancel(): void
 }
 
 export function isTaskQueue(value: any): boolean {

src/index.ts

@@ -83,6 +83,16 @@ class AbortError extends Error {
   }
 }
 
+class CancelError extends Error {
+  constructor() {
+    super('The task has been cancelled')
+  }
+
+  get name() {
+    return 'CancelError'
+  }
+}
+
 type ResourceLimits = Worker extends {
   resourceLimits?: infer T
 }
@@ -114,6 +124,13 @@ class ArrayTaskQueue implements TaskQueue {
     assert.notStrictEqual(index, -1)
     this.tasks.splice(index, 1)
   }
+
+  cancel(): void {
+    while (this.tasks.length > 0) {
+      const task = this.tasks.pop()
+      task?.cancel()
+    }
+  }
 }
 
 interface Options {
@@ -245,6 +262,7 @@ class TaskInfo extends AsyncResource implements Task {
   workerInfo: WorkerInfo | null = null
   created: number
   started: number
+  cancel: () => void
 
   constructor(
     task: any,
@@ -259,6 +277,7 @@ class TaskInfo extends AsyncResource implements Task {
     this.callback = callback
     this.task = task
     this.transferList = transferList
+    this.cancel = () => this.callback(new CancelError(), null)
 
     // If the task is a Transferable returned by
     // Tinypool.move(), then add it to the transferList
@@ -1072,6 +1091,11 @@ class Tinypool extends EventEmitterAsyncResource {
     return Math.max(pool.taskQueue.size - pool.pendingCapacity(), 0)
   }
 
+  cancelPendingTasks() {
+    const pool = this.#pool
+    pool.taskQueue.cancel()
+  }
+
   get completed(): number {
     return this.#pool.completed
   }

test/task-queue.test.ts

@@ -260,3 +260,51 @@ test('custom task queue works', async () => {
   expect(pushCalled).toBeTruthy()
   expect(shiftCalled).toBeTruthy()
 })
+
+test('queued tasks can be cancelled', async () => {
+  const pool = new Tinypool({
+    filename: resolve(__dirname, 'fixtures/sleep.js'),
+    minThreads: 0,
+    maxThreads: 1,
+  })
+
+  const time = 500
+  const taskCount = 10
+
+  const promises = []
+  let finishedTasks = 0
+  let cancelledTasks = 0
+
+  for (const _ of Array(taskCount)) {
+    const promise = pool
+      .run({ time })
+      .then(() => {
+        finishedTasks++
+      })
+      .catch((error) => {
+        if (error.message !== 'The task has been cancelled') {
+          throw error
+        }
+        cancelledTasks++
+      })
+    promises.push(promise)
+  }
+
+  // Wait for the first task to start
+  await new Promise((resolve) => setTimeout(resolve, time / 2))
+  expect(pool.queueSize).toBe(taskCount - 1)
+
+  // One task is running, cancel the pending ones
+  pool.cancelPendingTasks()
+
+  // The first task should still be on-going, pending ones should have started their cancellation
+  expect(finishedTasks).toBe(0)
+  expect(pool.queueSize).toBe(0)
+
+  await Promise.all(promises)
+
+  expect({ finishedTasks, cancelledTasks }).toEqual({
+    finishedTasks: 1,
+    cancelledTasks: taskCount - 1,
+  })
+})