feat: Custom Balancer

docs/docs/api-reference/event.md

@@ -27,4 +27,16 @@ by number of tasks enqueued that are pending of execution.
 
 ## Event: `'message'`
 
-A `'message'` event is emitted whenever a message is received from a worker thread.
+A `'message'` event is emitted whenever a message is received from a worker thread.
+
+## Event: `'workerCreate'`
+
+Event that is triggered when a new worker is created.
+
+As argument, it receives the worker instance.
+
+## Event: `'workerDestroy'`
+
+Event that is triggered when a worker is destroyed.
+
+As argument, it receives the worker instance that has been destroyed.

src/index.ts

@@ -18,7 +18,8 @@ import type {
 import {
   kQueueOptions,
   kTransferable,
-  kValue
+  kValue,
+  kWorkerData
 } from './symbols';
 import {
   TaskQueue,
@@ -32,7 +33,10 @@ import {
 } from './task_queue';
 import {
   WorkerInfo,
-  AsynchronouslyCreatedResourcePool
+  AsynchronouslyCreatedResourcePool,
+  PiscinaTaskBalancer,
+  PiscinaWorker,
+  ResourceBasedBalancer
 } from './worker_pool';
 import {
   AbortSignalAny,
@@ -71,7 +75,9 @@ interface Options {
   niceIncrement? : number,
   trackUnmanagedFds? : boolean,
   closeTimeout?: number,
-  recordTiming?: boolean
+  recordTiming?: boolean,
+  loadBalancer?: PiscinaTaskBalancer,
+  workerHistogram?: boolean,
 }
 
 interface FilledOptions extends Options {
@@ -86,7 +92,8 @@ interface FilledOptions extends Options {
   taskQueue : TaskQueue,
   niceIncrement : number,
   closeTimeout : number,
-  recordTiming : boolean
+  recordTiming : boolean,
+  workerHistogram: boolean,
 }
 
 interface RunOptions {
@@ -120,7 +127,8 @@ const kDefaultOptions : FilledOptions = {
   niceIncrement: 0,
   trackUnmanagedFds: true,
   closeTimeout: 30000,
-  recordTiming: true
+  recordTiming: true,
+  workerHistogram: false
 };
 
 const kDefaultRunOptions : FilledRunOptions = {
@@ -173,6 +181,7 @@ class ThreadPool {
   workerFailsDuringBootstrap : boolean = false;
   destroying : boolean = false;
   maxCapacity: number;
+  balancer: PiscinaTaskBalancer;
 
   constructor (publicInterface : Piscina, options : Options) {
     this.publicInterface = publicInterface;
@@ -202,9 +211,10 @@ class ThreadPool {
       this.options.maxQueue = options.maxQueue ?? kDefaultOptions.maxQueue;
     }
 
+    this.balancer = this.options.loadBalancer ?? ResourceBasedBalancer({ maximumUsage: this.options.concurrentTasksPerWorker });
     this.workers = new AsynchronouslyCreatedResourcePool<WorkerInfo>(
       this.options.concurrentTasksPerWorker);
-    this.workers.onAvailable((w : WorkerInfo) => this._onWorkerAvailable(w));
+    this.workers.onTaskDone((w : WorkerInfo) => this._onWorkerTaskDone(w));
     this.maxCapacity = this.options.maxThreads * this.options.concurrentTasksPerWorker;
 
     this.startingUp = true;
@@ -223,6 +233,8 @@ class ThreadPool {
   }
 
   _addNewWorker () : void {
+    if (this.closingUp) return;
+
     const pool = this;
     const worker = new Worker(resolve(__dirname, 'worker.js'), {
       env: this.options.env,
@@ -234,11 +246,27 @@ class ThreadPool {
     });
 
     const { port1, port2 } = new MessageChannel();
-    const workerInfo = new WorkerInfo(worker, port1, onMessage);
+    const workerInfo = new WorkerInfo(worker, port1, onMessage, this.options.workerHistogram);
+
+    workerInfo.onDestroy(() => {
+      this.publicInterface.emit('workerDestroy', workerInfo.interface);
+    });
+
     if (this.startingUp) {
       // There is no point in waiting for the initial set of Workers to indicate
       // that they are ready, we just mark them as such from the start.
       workerInfo.markAsReady();
+      // We need to emit the event in the next microtask, so that the user can
+      // attach event listeners before the event is emitted.
+      queueMicrotask(() => {
+        this.publicInterface.emit('workerCreate', workerInfo.interface);
+        this._onWorkerReady(workerInfo);
+      });
+    } else {
+      workerInfo.onReady(() => {
+        this.publicInterface.emit('workerCreate', workerInfo.interface);
+        this._onWorkerReady(workerInfo);
+      });
     }
 
     const message : StartupMessage = {
@@ -259,7 +287,9 @@ class ThreadPool {
       const taskInfo = workerInfo.taskInfos.get(taskId);
       workerInfo.taskInfos.delete(taskId);
 
-      pool.workers.maybeAvailable(workerInfo);
+      // TODO: we can abstract the task info handling
+      // right into the pool.workers.taskDone method
+      pool.workers.taskDone(workerInfo);
 
       /* istanbul ignore if */
       if (taskInfo === undefined) {
@@ -369,39 +399,81 @@ class ThreadPool {
     this.workers.delete(workerInfo);
   }
 
+  _onWorkerReady (workerInfo : WorkerInfo) : void {
+    this._onWorkerAvailable(workerInfo);
+  }
+
+  _onWorkerTaskDone (workerInfo: WorkerInfo) : void {
+    this._onWorkerAvailable(workerInfo);
+  }
+
   _onWorkerAvailable (workerInfo : WorkerInfo) : void {
-    while ((this.taskQueue.size > 0 || this.skipQueue.length > 0) &&
-      workerInfo.currentUsage() < this.options.concurrentTasksPerWorker) {
+    let workers: PiscinaWorker[] | null = null;
+    while ((this.taskQueue.size > 0 || this.skipQueue.length > 0)) {
       // The skipQueue will have tasks that we previously shifted off
       // the task queue but had to skip over... we have to make sure
       // we drain that before we drain the taskQueue.
       const taskInfo = this.skipQueue.shift() ||
                        this.taskQueue.shift() as TaskInfo;
-      // If the task has an abortSignal and the worker has any other
-      // tasks, we cannot distribute the task to it. Skip for now.
-      if (taskInfo.abortSignal && workerInfo.taskInfos.size > 0) {
-        this.skipQueue.push(taskInfo);
-        break;
+
+      if (workers == null) {
+        workers = [...this.workers].map(workerInfo => workerInfo.interface);
       }
-      const now = performance.now();
-      this.waitTime?.record(toHistogramIntegerNano(now - taskInfo.created));
-      taskInfo.started = now;
-      workerInfo.postTask(taskInfo);
-      this._maybeDrain();
-      return;
+
+      const distributed = this._distributeTask(taskInfo, workers);
+
+      // If task was distributed, we should continue to distribute more tasks
+      if (distributed) {
+        continue;
+        // If balancer states that pool is busy, we should stop trying to distribute tasks
+      } else { break; }
     }
 
-    if (workerInfo.taskInfos.size === 0 &&
+    // If more workers than minThreads, we can remove idle workers
+    if (workerInfo.currentUsage() === 0 &&
         this.workers.size > this.options.minThreads) {
       workerInfo.idleTimeout = setTimeout(() => {
-        assert.strictEqual(workerInfo.taskInfos.size, 0);
+        assert.strictEqual(workerInfo.currentUsage(), 0);
         if (this.workers.size > this.options.minThreads) {
           this._removeWorker(workerInfo);
         }
       }, this.options.idleTimeout).unref();
     }
   }
 
+  _distributeTask (task: TaskInfo, workers: PiscinaWorker[]): boolean {
+    // We need to verify if the task is aborted already or not
+    // otherwise we might be distributing aborted tasks to workers
+    if (task.aborted) return false;
+
+    const candidate = this.balancer(task.interface, workers);
+
+    // Seeking for a real worker instead of customized one
+    if (candidate != null && candidate[kWorkerData] != null) {
+      const now = performance.now();
+      this.waitTime?.record(toHistogramIntegerNano(now - task.created));
+      task.started = now;
+      candidate[kWorkerData].postTask(task);
+      this._maybeDrain();
+      // If candidate, let's try to distribute more tasks
+      return true;
+    }
+
+    // We spawn if possible
+    // TODO: scheduler will intercept this.
+    if (this.workers.size < this.options.maxThreads) {
+      this._addNewWorker();
+    }
+
+    if (task.abortSignal) {
+      this.skipQueue.push(task);
+    } else {
+      this.taskQueue.push(task);
+    }
+
+    return false;
+  }
+
   runTask (
     task : any,
     options : RunOptions) : Promise<any> {
@@ -424,9 +496,9 @@ class ThreadPool {
     filename = maybeFileURLToPath(filename);
 
     let signal: AbortSignalAny | null;
-    if (this.closingUp) {
+    if (this.closingUp || this.destroying) {
       const closingUpAbortController = new AbortController();
-      closingUpAbortController.abort('queue is closing up');
+      closingUpAbortController.abort('queue is being terminated');
 
       signal = closingUpAbortController.signal;
     } else {
@@ -462,8 +534,10 @@ class ThreadPool {
       // If the AbortSignal has an aborted property and it's truthy,
       // reject immediately.
       if ((signal as AbortSignalEventTarget).aborted) {
-        return Promise.reject(new AbortError((signal as AbortSignalEventTarget).reason));
+        reject!(new AbortError((signal as AbortSignalEventTarget).reason));
+        return ret;
       }
+
       taskInfo.abortListener = () => {
         // Call reject() first to make sure we always reject with the AbortError
         // if the task is aborted, not with an Error from the possible
@@ -476,68 +550,40 @@ class ThreadPool {
           this._ensureMinimumWorkers();
         } else {
           // Not yet running: Remove it from the queue.
+          // Call should be idempotent
           this.taskQueue.remove(taskInfo);
         }
       };
+
       onabort(signal, taskInfo.abortListener);
     }
 
-    // If there is a task queue, there's no point in looking for an available
-    // Worker thread. Add this task to the queue, if possible.
     if (this.taskQueue.size > 0) {
       const totalCapacity = this.options.maxQueue + this.pendingCapacity();
       if (this.taskQueue.size >= totalCapacity) {
         if (this.options.maxQueue === 0) {
-          return Promise.reject(Errors.NoTaskQueueAvailable());
+          reject!(Errors.NoTaskQueueAvailable());
         } else {
-          return Promise.reject(Errors.TaskQueueAtLimit());
+          reject!(Errors.TaskQueueAtLimit());
         }
       } else {
-        if (this.workers.size < this.options.maxThreads) {
-          this._addNewWorker();
-        }
         this.taskQueue.push(taskInfo);
       }
 
       this._maybeDrain();
       return ret;
     }
 
-    // Look for a Worker with a minimum number of tasks it is currently running.
-    let workerInfo : WorkerInfo | null = this.workers.findAvailable();
-
-    // If we want the ability to abort this task, use only workers that have
-    // no running tasks.
-    if (workerInfo !== null && workerInfo.currentUsage() > 0 && signal) {
-      workerInfo = null;
-    }
-
-    // If no Worker was found, or that Worker was handling another task in some
-    // way, and we still have the ability to spawn new threads, do so.
-    let waitingForNewWorker = false;
-    if ((workerInfo === null || workerInfo.currentUsage() > 0) &&
-        this.workers.size < this.options.maxThreads) {
-      this._addNewWorker();
-      waitingForNewWorker = true;
-    }
+    const workers = [...this.workers.readyItems].map(workerInfo => workerInfo.interface);
+    const distributed = this._distributeTask(taskInfo, workers);
 
-    // If no Worker is found, try to put the task into the queue.
-    if (workerInfo === null) {
-      if (this.options.maxQueue <= 0 && !waitingForNewWorker) {
-        return Promise.reject(Errors.NoTaskQueueAvailable());
-      } else {
-        this.taskQueue.push(taskInfo);
+    if (!distributed) {
+      // We reject if no task queue set and no more pending capacity.
+      if (this.options.maxQueue <= 0 && this.pendingCapacity() === 0) {
+        reject!(Errors.NoTaskQueueAvailable());
       }
+    };
 
-      this._maybeDrain();
-      return ret;
-    }
-
-    // TODO(addaleax): Clean up the waitTime/runTime recording.
-    const now = performance.now();
-    this.waitTime?.record(toHistogramIntegerNano(now - taskInfo.created));
-    taskInfo.started = now;
-    workerInfo.postTask(taskInfo);
     this._maybeDrain();
     return ret;
   }
@@ -639,12 +685,12 @@ class ThreadPool {
       for (const workerInfo of this.workers) {
         checkIfWorkerIsDone(workerInfo);
 
-        workerInfo.port.on('message', () => checkIfWorkerIsDone(workerInfo));
+        this.workers.onTaskDone(checkIfWorkerIsDone);
       }
     });
 
     const throwOnTimeOut = async (timeout: number) => {
-      await sleep(timeout);
+      await sleep(timeout, null, { ref: false });
       throw Errors.CloseTimeout();
     };
 
@@ -725,6 +771,12 @@ export default class Piscina<T = any, R = any> extends EventEmitterAsyncResource
     if (options.closeTimeout !== undefined && (typeof options.closeTimeout !== 'number' || options.closeTimeout < 0)) {
       throw new TypeError('options.closeTimeout must be a non-negative integer');
     }
+    if (options.loadBalancer !== undefined && (typeof options.loadBalancer !== 'function' || options.loadBalancer.length < 1)) {
+      throw new TypeError('options.loadBalancer must be a function with at least two args');
+    }
+    if (options.workerHistogram !== undefined && (typeof options.workerHistogram !== 'boolean')) {
+      throw new TypeError('options.workerHistogram must be a boolean');
+    }
 
     this.#pool = new ThreadPool(this, options);
   }
@@ -755,6 +807,7 @@ export default class Piscina<T = any, R = any> extends EventEmitterAsyncResource
       return Promise.reject(
         new TypeError('signal argument must be an object'));
     }
+
     return this.#pool.runTask(task, { transferList, filename, name, signal });
   }
 
@@ -829,7 +882,7 @@ export default class Piscina<T = any, R = any> extends EventEmitterAsyncResource
     }
 
     // count is available as of Node.js v16.14.0 but not present in the types
-    const count = (this.#pool.runTime as RecordableHistogram & { count: number}).count;
+    const count = (this.#pool.runTime as RecordableHistogram & { count: number }).count;
     if (count === 0) {
       return 0;
     }

src/symbols.ts

@@ -1,6 +1,7 @@
 // Internal symbol used to mark Transferable objects returned
 // by the Piscina.move() function
 export const kMovable = Symbol('Piscina.kMovable');
+export const kWorkerData = Symbol('Piscina.kWorkerData');
 export const kTransferable = Symbol.for('Piscina.transferable');
 export const kValue = Symbol.for('Piscina.valueOf');
 export const kQueueOptions = Symbol.for('Piscina.queueOptions');