Merge branch 'current' into typescript-entrypoint

Author: metcoder95

docs/docs/api-reference/class.md

@@ -223,7 +223,7 @@ If the `PiscinaLoadBalancer` returns `null`, `Piscina` will attempt to spawn a n
 
 ```ts
 interface PiscinaTask {
-  taskId: number; // Unique identifier for the task
+  taskId: string; // Unique identifier for the task
   filename: string; // Filename of the worker module
   name: string; // Name of the worker function
   created: number; // Timestamp when the task was created

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/common.ts

@@ -9,7 +9,7 @@ export interface TaskQueue {
 
 // Public Interface
 export interface PiscinaTask extends Task {
-    taskId: number;
+    taskId: string;
     filename: string;
     name: string;
     created: number;

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
  *
@@ -41,32 +51,43 @@ export function isTaskQueue (value: TaskQueue): boolean {
   );
 }
 
-let taskIdCounter = 0;
+
+function taskIdFactory() {
+  let taskIdCounter = 0;
+  const maxint = 2147483647;
+  return () => {
+    taskIdCounter = (taskIdCounter + 1) & maxint;
+    return taskIdCounter.toString(36);
+  }
+}
+
 // Extend AsyncResource so that async relations between posting a task and
 // receiving its result are visible to diagnostic tools.
 export class TaskInfo extends AsyncResource implements Task {
+    static getTaskId: () => string = taskIdFactory();
     callback : TaskCallback;
     task : any;
     transferList : TransferList;
     filename : string;
     name : string;
-    taskId : number;
+    taskId : string;
     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,34 +99,31 @@ 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];
       }
 
       this.filename = filename;
       this.name = name;
-      // TODO: This should not be global
-      this.taskId = taskIdCounter++;
+      this.taskId = TaskInfo.getTaskId();
       this.abortSignal = abortSignal;
       this.created = performance.now();
       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 +132,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 {

src/types.ts

@@ -13,7 +13,7 @@ export interface StartupMessage {
 }
 
 export interface RequestMessage {
-  taskId: number
+  taskId: string
   task: any
   filename: string
   name: string
@@ -25,7 +25,7 @@ export interface ReadyMessage {
 }
 
 export interface ResponseMessage {
-  taskId: number
+  taskId: string
   result: any
   error: Error | null
   time: number | null

src/worker_pool/index.ts

@@ -33,7 +33,7 @@ type WorkerInfoParams = {
 
 export class WorkerInfo extends AsynchronouslyCreatedResource {
     worker : Worker;
-    taskInfos : Map<number, TaskInfo>;
+    taskInfos : Map<string, TaskInfo>;
     idleTimeout : NodeJS.Timeout | null = null;
     port : MessagePort;
     sharedBuffer : Int32Array;
@@ -184,7 +184,7 @@ export class WorkerInfo extends AsynchronouslyCreatedResource {
       return this.taskInfos.size;
     }
 
-    popTask (taskId: number) : TaskInfo | null {
+    popTask (taskId: string) : TaskInfo | null {
       const task = this.taskInfos.get(taskId) ?? null;
 
       this.taskInfos.delete(taskId);