wip: fixes and adding transferable support

[ci skip]

Author: tegefaulkes

src/WorkerManager.ts

@@ -63,13 +63,13 @@ class WorkerManager {
   }
 
   @ready(new errors.ErrorWorkerManagerDestroyed())
-  public async call<T>(task: WorkerTask): Promise<T> {
+  public async call(task: WorkerTask): Promise<unknown> {
     return await this.queue(task);
   }
 
   @ready(new errors.ErrorWorkerManagerDestroyed())
-  public queue<T>(task: WorkerTask): Promise<T> {
-    return new Promise<T>((resolve, reject) => {
+  public queue(task: WorkerTask): Promise<unknown> {
+    return new Promise((resolve, reject) => {
       this.pool.runTask(task, (result, error) => {
         if (error != null) return reject(error);
         return resolve(result);

src/WorkerPool.ts

@@ -1,10 +1,17 @@
 import type { Worker } from 'node:worker_threads';
 import type { Subscription } from 'rxjs';
-import type { WorkerFactory, WorkerTask } from './types.js';
+import type {
+  TaskCallback,
+  WorkerFactory,
+  WorkerResultInternal,
+  WorkerTask,
+  WorkerTaskInput,
+} from './types.js';
 import { Subject } from 'rxjs';
+import { WorkerFunction, WorkerResult } from './types.js';
+import * as errors from './errors.js';
 
 const taskInfoSymbol = Symbol('Task Info Symbol');
-type TaskCallback = (result: any | undefined, error: Error | undefined) => void;
 type PoolStatus = 'idle' | 'working' | 'queued';
 
 class WorkerPool {
@@ -23,7 +30,7 @@ class WorkerPool {
   public poolStatus: PoolStatus = 'idle';
 
   constructor(workerNum: number, workerFactory: WorkerFactory) {
-    if (workerNum < 1) throw Error('TMP IMP must be at least 1 worker');
+    if (workerNum < 1) throw new errors.ErrorWorkerPoolInvalidWorkers();
     this.workerFactory = workerFactory;
     for (let i = 0; i < workerNum; i++) {
       this.addWorker();
@@ -50,9 +57,9 @@ class WorkerPool {
 
   protected addWorker() {
     const worker = this.workerFactory();
-    const messageHandler = (result: { result?: unknown; error?: Error }) => {
+    const messageHandler = (result: WorkerResultInternal) => {
       if (result.error != null) worker[taskInfoSymbol](undefined, result.error);
-      else worker[taskInfoSymbol](result.result);
+      else worker[taskInfoSymbol](result.data, undefined);
       worker[taskInfoSymbol] = undefined;
       this.freeWorkers.push(worker);
       this.$workerFreed.next();
@@ -83,7 +90,7 @@ class WorkerPool {
     this.$workerFreed.next();
   }
 
-  public runTask(task: any, callback: any) {
+  public runTask(task: WorkerTaskInput, callback: TaskCallback) {
     if (this.terminating) throw Error('TMP IMP terminating');
     if (this.freeWorkers.length === 0) {
       this.queue.push({ task, callback });
@@ -95,33 +102,13 @@ class WorkerPool {
     const worker = this.freeWorkers.pop()!;
     if (wasIdle) this.$poolStatus.next('working');
     worker[taskInfoSymbol] = callback;
-    worker.postMessage(task);
+    worker.postMessage({ type: task.type, data: task.data }, task.transferList);
   }
 
   public async terminate(force: boolean) {
     this.terminating = true;
-    if (!force) {
-      // Prevent new tasks and wait for exising queue to drain
-      await new Promise<void>((resolve) => {
-        if (
-          this.freeWorkers.length === this.workers.size &&
-          this.queue.length === 0
-        ) {
-          return resolve();
-        }
-        const subscription = this.$workerFreed.subscribe(() => {
-          // Wait for the queue to drain and the workers to free
-          if (
-            this.freeWorkers.length === this.workers.size &&
-            this.queue.length === 0
-          ) {
-            subscription.unsubscribe();
-            resolve();
-          }
-        });
-      });
-    }
-
+    // Prevent new tasks and wait for exising queue to drain
+    if (!force) await this.settled();
     // Prevent terminations from creating new workers
     this.handleDestroySubscription.unsubscribe();
     const workerTerminatePs: Array<Promise<number>> = [];

src/errors.ts

@@ -8,6 +8,14 @@ class ErrorWorkerManagerDestroyed<T> extends ErrorWorkerManager<T> {
   static description = 'WorkerManager is destroyed';
 }
 
+class ErrorWorkerPool<T> extends AbstractError<T> {
+  static description = 'WorkerPool error';
+}
+
+class ErrorWorkerPoolInvalidWorkers<T> extends ErrorWorkerManager<T> {
+  static description = 'You must have at minimum 1 worker';
+}
+
 class ErrorWorker<T> extends AbstractError<T> {
   static description = 'Worker error';
 }
@@ -19,6 +27,8 @@ class ErrorWorkerHandlerMissing<T> extends ErrorWorkerManager<T> {
 export {
   ErrorWorkerManager,
   ErrorWorkerManagerDestroyed,
+  ErrorWorkerPool,
+  ErrorWorkerPoolInvalidWorkers,
   ErrorWorker,
   ErrorWorkerHandlerMissing,
 };

src/expose.ts

@@ -1,4 +1,9 @@
-import type { WorkerTask, WorkerManifest } from '#types.js';
+import type {
+  WorkerTask,
+  WorkerManifest,
+  WorkerResult,
+  WorkerResultInternal,
+} from '#types.js';
 import { isMainThread, parentPort, threadId } from 'node:worker_threads';
 import * as workerErrors from '#errors.js';
 
@@ -18,14 +23,14 @@ function expose(workerManifest: WorkerManifest) {
       });
       return;
     }
-    let result: unknown;
+    let result: WorkerResultInternal;
     try {
       result = await method(task.data);
     } catch (error) {
       parentPort!.postMessage({ error });
       return;
     }
-    parentPort!.postMessage({ result });
+    parentPort!.postMessage({ data: result.data }, result.transferList);
   };
   const handleMessageError = (e_) => {
     // Do nothing for now

src/types.ts

@@ -1,13 +1,35 @@
-import type { Worker } from 'node:worker_threads';
+import type { TransferListItem, Worker } from 'node:worker_threads';
 
 type WorkerFactory = () => Worker;
 
-type WorkerTask<T extends string = string, D extends any = any> = {
+type WorkerTask<T extends string = string, D = unknown> = {
   type: T;
   data: D;
 };
 
-type WorkerFunction<I = any, O = any> = (data: I) => Promise<O>;
+type WorkerTaskInput = WorkerTask & {
+  transferList?: Array<TransferListItem>;
+};
+
+type WorkerFunction<I = unknown, O = unknown> = (
+  data: I,
+) => Promise<WorkerResult<O>>;
+type WorkerResult<O = unknown> = { data: O };
+type WorkerResultInternal<O = unknown> = WorkerResult<O> & {
+  error?: Error;
+  transferList?: Array<TransferListItem>;
+};
 type WorkerManifest = Record<string, WorkerFunction>;
 
-export type { WorkerFactory, WorkerTask, WorkerFunction, WorkerManifest };
+type TaskCallback<O = unknown> = (result: O, error?: Error) => void;
+
+export type {
+  WorkerFactory,
+  WorkerTask,
+  WorkerTaskInput,
+  WorkerFunction,
+  WorkerResult,
+  WorkerResultInternal,
+  WorkerManifest,
+  TaskCallback,
+};

src/worker.ts

@@ -5,27 +5,28 @@ import { expose } from './expose.js';
 
 const worker = {
   test: async () => {
-    return 'hello world!';
+    return { data: 'hello world!' };
   },
-  add: async (data: { a: number; b: number }): Promise<number> => {
-    return data.a + data.b;
+  add: async (data: { a: number; b: number }) => {
+    return { data: data.a + data.b };
   },
-  sub: async (data: { a: number; b: number }): Promise<number> => {
-    return data.a - data.b;
+  sub: async (data: { a: number; b: number }) => {
+    return { data: data.a - data.b };
   },
-  fac: async (data: number): Promise<number> => {
+  fac: async (data: number) => {
     let acc = 1;
     for (let i = 1; i <= data; i++) {
       acc = acc * i;
     }
-    return acc;
+    return { data: acc };
   },
   sleep: async (data: number) => {
     await new Promise<void>((resolve) => {
       const timer = setTimeout(resolve, data);
       // Prevent holding the process open
       timer.unref();
     });
+    return { data: undefined };
   },
 } satisfies WorkerManifest;