refactor: TaskInfo (#857)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
(cherry picked from commit ec213ad)

Author: 2 people

src/abort.ts

@@ -1,3 +1,5 @@
+import type { EventEmitter } from 'node:events';
+
 interface AbortSignalEventTargetAddOptions {
   once: boolean;
 }
@@ -32,10 +34,12 @@ export class AbortError extends Error {
   }
 }
 
-export function onabort (abortSignal: AbortSignalAny, listener: () => void) {
+export function onabort (abortSignal: AbortSignalAny, listener: () => void): () => void {
   if ('addEventListener' in abortSignal) {
     abortSignal.addEventListener('abort', listener, { once: true });
+    return () => abortSignal.removeEventListener('abort', listener);
   } else {
     abortSignal.once('abort', listener);
+    return () => (abortSignal as EventEmitter).removeListener('abort', listener);
   }
 }

src/index.ts

@@ -40,7 +40,6 @@ import {
   AbortSignalAny,
   AbortSignalEventTarget,
   AbortError,
-  onabort
 } from './abort';
 import {
   PiscinaHistogram,
@@ -507,12 +506,15 @@ class ThreadPool {
     }
 
     const { promise: ret, resolve, reject } = promiseResolvers();
-    const taskInfo = new TaskInfo(
+    const taskInfo = new TaskInfo({
       task,
       transferList,
       filename,
       name,
-      (err : Error | null, result : any) => {
+      abortSignal: signal,
+      triggerAsyncId: this.publicInterface.asyncResource.asyncId()
+    },
+    (err : Error | null, result : any) => {
         this.completed++;
         if (taskInfo.started) {
           this.histogram?.recordRunTime(performance.now() - taskInfo.started);
@@ -524,9 +526,7 @@ class ThreadPool {
         }
 
         queueMicrotask(this._maybeDrain.bind(this))
-      },
-      signal,
-      this.publicInterface.asyncResource.asyncId());
+    });
 
     if (signal != null) {
       // If the AbortSignal has an aborted property and it's truthy,
@@ -536,7 +536,7 @@ class ThreadPool {
         return ret;
       }
 
-      taskInfo.abortListener = () => {
+      taskInfo.onAbort(() => {
         // Call reject() first to make sure we always reject with the AbortError
         // if the task is aborted, not with an Error from the possible
         // thread termination below.
@@ -551,9 +551,9 @@ class ThreadPool {
           // Call should be idempotent
           this.taskQueue.remove(taskInfo);
         }
-      };
+      });
 
-      onabort(signal, taskInfo.abortListener);
+      taskInfo.setAbortListener(signal);
     }
 
     if (this.taskQueue.size > 0) {

src/task_queue/index.ts

@@ -3,11 +3,12 @@ import { performance } from 'node:perf_hooks';
 import { AsyncResource } from 'node:async_hooks';
 
 import type { WorkerInfo } from '../worker_pool';
-import type { AbortSignalAny, AbortSignalEventEmitter } from '../abort';
+import type { Task, TaskQueue, PiscinaTask } from './common';
+
+import { onabort, type AbortSignalAny } from '../abort';
 import { isMovable } from '../common';
 import { kTransferable, kValue, kQueueOptions } from '../symbols';
 
-import type { Task, TaskQueue, PiscinaTask } from './common';
 
 export { ArrayTaskQueue } from './array_queue';
 export { FixedQueue } from './fixed_queue';
@@ -23,6 +24,15 @@ export type TransferList = MessagePort extends {
   : never
 export type TransferListItem = TransferList extends Array<infer T> ? T : never
 
+type TaskInfoParameters = {
+  task : any;
+  transferList : TransferList;
+  filename : string;
+  name : string;
+  abortSignal : AbortSignalAny | null;
+  triggerAsyncId : number;
+}
+
 /**
  * Verifies if a given TaskQueue is valid
  *
@@ -52,21 +62,22 @@ export class TaskInfo extends AsyncResource implements Task {
     name : string;
     taskId : number;
     abortSignal : AbortSignalAny | null;
-    // abortListener : (() => void) | null = null;
     workerInfo : WorkerInfo | null = null;
     created : number;
     started : number;
     aborted = false;
-    _abortListener: (() => void) | null = null;
-
-    constructor (
-      task : any,
-      transferList : TransferList,
-      filename : string,
-      name : string,
-      callback : TaskCallback,
-      abortSignal : AbortSignalAny | null,
-      triggerAsyncId : number) {
+    _abortListener: (() => void) = () => { this.aborted = true; };
+    _abortCleaner: (() => void) | null = null;
+
+    constructor ({
+      task,
+      transferList,
+      filename,
+      name,
+      abortSignal,
+      triggerAsyncId,
+    }: TaskInfoParameters,
+    callback: TaskCallback) {
       super('Piscina.Task', { requireManualDestroy: true, triggerAsyncId });
       this.callback = callback;
       this.task = task;
@@ -78,12 +89,11 @@ export class TaskInfo extends AsyncResource implements Task {
       if (isMovable(task)) {
         // This condition should never be hit but typescript
         // complains if we dont do the check.
-        /* istanbul ignore if */
+        /* c8 ignore next */
         if (this.transferList == null) {
           this.transferList = [];
         }
-        this.transferList =
-          this.transferList.concat(task[kTransferable]);
+        this.transferList = this.transferList.concat(task[kTransferable]);
         this.task = task[kValue];
       }
 
@@ -96,16 +106,15 @@ export class TaskInfo extends AsyncResource implements Task {
       this.started = 0;
     }
 
-    // TODO: improve this handling - ideally should be extended
-    set abortListener (value: (() => void)) {
+    onAbort (value: (() => void)) {
       this._abortListener = () => {
         this.aborted = true;
         value();
       };
     }
 
-    get abortListener (): (() => void) | null {
-      return this._abortListener;
+    setAbortListener(signal: AbortSignalAny) : void {
+      this._abortCleaner = onabort(signal, this._abortListener);
     }
 
     releaseTask () : any {
@@ -114,19 +123,15 @@ export class TaskInfo extends AsyncResource implements Task {
       return ret;
     }
 
+    // TODO: implement - helpful for streaming chunks of data from worker to parent
+    onResponse(_result: any) {}
+
     done (err : Error | null, result? : any) : void {
       this.runInAsyncScope(this.callback, null, err, result);
       this.emitDestroy(); // `TaskInfo`s are used only once.
       // If an abort signal was used, remove the listener from it when
       // done to make sure we do not accidentally leak.
-      if (this.abortSignal && this.abortListener) {
-        if ('removeEventListener' in this.abortSignal && this.abortListener) {
-          this.abortSignal.removeEventListener('abort', this.abortListener);
-        } else {
-          (this.abortSignal as AbortSignalEventEmitter).off(
-            'abort', this.abortListener);
-        }
-      }
+      this._abortCleaner?.();
     }
 
     get [kQueueOptions] () : {} | null {