wip: handling workers that fail to create properly and test expanding

Author: tegefaulkes

src/WorkerManager.ts

@@ -1,7 +1,7 @@
 import type { WorkerFactory } from './types.js';
 import Logger from '@matrixai/logger';
 import { CreateDestroy, ready } from '@matrixai/async-init/CreateDestroy.js';
-import { WorkerTask } from './types.js';
+import { WorkerTaskInformation } from './types.js';
 import * as errors from './errors.js';
 import WorkerPool from './WorkerPool.js';
 
@@ -63,12 +63,12 @@ class WorkerManager {
   }
 
   @ready(new errors.ErrorWorkerManagerDestroyed())
-  public async call(task: WorkerTask): Promise<unknown> {
+  public async call(task: WorkerTaskInformation): Promise<unknown> {
     return await this.queue(task);
   }
 
   @ready(new errors.ErrorWorkerManagerDestroyed())
-  public queue(task: WorkerTask): Promise<unknown> {
+  public queue(task: WorkerTaskInformation): Promise<unknown> {
     return new Promise((resolve, reject) => {
       this.pool.runTask(task, (result, error) => {
         if (error != null) return reject(error);

src/WorkerPool.ts

@@ -4,22 +4,40 @@ import type {
   TaskCallback,
   WorkerFactory,
   WorkerResultInternal,
-  WorkerTask,
+  WorkerTaskInformation,
   WorkerTaskInput,
 } from './types.js';
+import { AsyncResource } from 'node:async_hooks';
 import { Subject } from 'rxjs';
 import { WorkerFunction, WorkerResult } from './types.js';
 import * as errors from './errors.js';
 
+class WorkerTask extends AsyncResource {
+  protected callback: TaskCallback;
+
+  constructor(callback: TaskCallback, triggerAsyncId?: number) {
+    super('WorkerTask', triggerAsyncId);
+    this.callback = callback;
+  }
+
+  public done(result, error: Error) {
+    this.runInAsyncScope(this.callback, null, result, error);
+    this.emitDestroy();
+  }
+}
+
 const taskInfoSymbol = Symbol('Task Info Symbol');
-type PoolStatus = 'idle' | 'working' | 'queued';
+type PoolStatus = 'terminated' | 'idle' | 'working' | 'queued';
 
 class WorkerPool {
   protected workerFactory: WorkerFactory;
   protected workers: Set<Worker> = new Set();
   protected freeWorkers: Array<Worker> = [];
-  protected queue: Array<{ task: WorkerTask; callback: TaskCallback }> = [];
-  protected terminating: boolean = false;
+  protected queue: Array<{
+    task: WorkerTaskInformation;
+    callback: TaskCallback;
+  }> = [];
+  protected terminatedError: Error | undefined = undefined;
   protected handleDestroySubscription: Subscription;
 
   public $workerCreated = new Subject<void>();
@@ -57,15 +75,31 @@ class WorkerPool {
 
   protected addWorker() {
     const worker = this.workerFactory();
+    let workerError: Error;
+    let initializing = true;
     const messageHandler = (result: WorkerResultInternal) => {
-      if (result.error != null) worker[taskInfoSymbol](undefined, result.error);
-      else worker[taskInfoSymbol](result.data, undefined);
+      if (initializing) {
+        // @ts-ignore: ignoring type here
+        if (result !== 'initialized') {
+          throw Error('TMP IMP failed to initialize properly');
+        }
+        initializing = false;
+        this.freeWorkers.push(worker);
+        this.$workerFreed.next();
+        return;
+      }
+      if (result.error != null) {
+        worker[taskInfoSymbol].done(undefined, result.error);
+      } else {
+        worker[taskInfoSymbol].done(result.data, undefined);
+      }
       worker[taskInfoSymbol] = undefined;
       this.freeWorkers.push(worker);
       this.$workerFreed.next();
     };
     const errorHandler = (e) => {
-      if (worker[taskInfoSymbol]) worker[taskInfoSymbol](undefined, e);
+      workerError = e;
+      if (worker[taskInfoSymbol]) worker[taskInfoSymbol].done(undefined, e);
       else this.$workerError.next(e);
     };
     worker.on('message', messageHandler);
@@ -74,24 +108,37 @@ class WorkerPool {
       worker.off('message', messageHandler);
       worker.off('error', errorHandler);
       this.workers.delete(worker);
+      if (
+        workerError != null &&
+        `${workerError.message}`.includes('Cannot find module')
+      ) {
+        // If the worker errored then we want to check if it was a failure to load error
+        if (this.workers.size === 0 && this.terminatedError == null) {
+          this.terminatedError = new Error(
+            'TMP IMP Workers failed to load modules',
+          );
+          this.cleanUp(this.terminatedError);
+        }
+        return;
+      }
       this.$workerDestroyed.next();
     });
+    worker.once('online', async () => {
+      worker.postMessage({ type: 'initialize', data: undefined });
+    });
 
     // TODO: debugging
     // worker.on('message', (...args) => console.log('DEBUG message: ', args));
     // worker.on('messageerror', (...args) => console.log('DEBUG messageerror: ', args));
     // worker.on('error', (...args) => console.log('DEBUG error: ', args));
     // worker.on('exit', (...args) => console.log('DEBUG exit: ', args));
     // worker.on('online', (...args) => console.log('DEBUG online: ', args));
-
     this.workers.add(worker);
     this.$workerCreated.next();
-    this.freeWorkers.push(worker);
-    this.$workerFreed.next();
   }
 
   public runTask(task: WorkerTaskInput, callback: TaskCallback) {
-    if (this.terminating) throw Error('TMP IMP terminating');
+    if (this.terminatedError != null) throw this.terminatedError;
     if (this.freeWorkers.length === 0) {
       this.queue.push({ task, callback });
       if (this.queue.length === 1) this.$poolStatus.next('queued');
@@ -101,12 +148,13 @@ class WorkerPool {
     const wasIdle = this.freeWorkers.length === this.workers.size;
     const worker = this.freeWorkers.pop()!;
     if (wasIdle) this.$poolStatus.next('working');
-    worker[taskInfoSymbol] = callback;
+    worker[taskInfoSymbol] = new WorkerTask(callback);
     worker.postMessage({ type: task.type, data: task.data }, task.transferList);
   }
 
   public async terminate(force: boolean) {
-    this.terminating = true;
+    if (this.terminatedError != null) return;
+    this.terminatedError = Error('TMP IMP terminating');
     // Prevent new tasks and wait for exising queue to drain
     if (!force) await this.settled();
     // Prevent terminations from creating new workers
@@ -116,9 +164,18 @@ class WorkerPool {
       workerTerminatePs.push(worker.terminate());
     }
     await Promise.all(workerTerminatePs);
-    // Console.log('workers', this.workers.size);
-    // console.log('queue', this.queue.length);
-    // console.log('freeWorkers', this.freeWorkers.length);
+    this.cleanUp(this.terminatedError);
+  }
+
+  protected cleanUp(terminatedError: Error): void {
+    // Cleaning up remaining queue and observables
+    let task = this.queue.pop();
+    while (task != null) {
+      const workerTask = new WorkerTask(task.callback);
+      workerTask.done(undefined, terminatedError);
+      task = this.queue.pop();
+    }
+    this.$poolStatus.next('terminated');
     // Cleaning up subjects
     this.$workerCreated.complete();
     this.$workerDestroyed.complete();
@@ -133,14 +190,16 @@ class WorkerPool {
    */
   public completed(): Promise<void> {
     return new Promise<void>((resolve, reject) => {
-      if (this.poolStatus === 'idle') return resolve();
+      if (this.poolStatus === 'idle' || this.poolStatus === 'terminated') {
+        return resolve();
+      }
       const errorSubscription = this.$workerError.subscribe((e) => {
         errorSubscription.unsubscribe();
         stateSubscription.unsubscribe();
         reject(e);
       });
       const stateSubscription = this.$poolStatus.subscribe((v) => {
-        if (v === 'idle') {
+        if (v === 'idle' || this.poolStatus === 'terminated') {
           errorSubscription.unsubscribe();
           stateSubscription.unsubscribe();
           return resolve();
@@ -154,9 +213,11 @@ class WorkerPool {
    */
   public settled(): Promise<void> {
     return new Promise<void>((resolve) => {
-      if (this.poolStatus === 'idle') return resolve();
+      if (this.poolStatus === 'idle' || this.poolStatus === 'terminated') {
+        return resolve();
+      }
       const subscription = this.$poolStatus.subscribe((v) => {
-        if (v === 'idle') {
+        if (v === 'idle' || this.poolStatus === 'terminated') {
           subscription.unsubscribe();
           return resolve();
         }

src/expose.ts

@@ -1,5 +1,5 @@
 import type {
-  WorkerTask,
+  WorkerTaskInformation,
   WorkerManifest,
   WorkerResult,
   WorkerResultInternal,
@@ -15,7 +15,11 @@ import * as workerErrors from '#errors.js';
 function expose(workerManifest: WorkerManifest) {
   // We don't want to run this in the main thread since it's not acting as a worker
   if (isMainThread) return;
-  const handleMessage = async (task: WorkerTask) => {
+  const handleMessage = async (task: WorkerTaskInformation) => {
+    if (task.type === 'initialize') {
+      parentPort!.postMessage('initialized');
+      return;
+    }
     const method = workerManifest[task.type];
     if (method == null) {
       parentPort!.postMessage({

src/types.ts

@@ -2,12 +2,12 @@ import type { TransferListItem, Worker } from 'node:worker_threads';
 
 type WorkerFactory = () => Worker;
 
-type WorkerTask<T extends string = string, D = unknown> = {
+type WorkerTaskInformation<T extends string = string, D = unknown> = {
   type: T;
   data: D;
 };
 
-type WorkerTaskInput = WorkerTask & {
+type WorkerTaskInput = WorkerTaskInformation & {
   transferList?: Array<TransferListItem>;
 };
 
@@ -25,7 +25,7 @@ type TaskCallback<O = unknown> = (result: O, error?: Error) => void;
 
 export type {
   WorkerFactory,
-  WorkerTask,
+  WorkerTaskInformation,
   WorkerTaskInput,
   WorkerFunction,
   WorkerResult,