feat!: drain/needsDrain accounts for task queue (#872)

Author: metcoder95

docs/docs/api-reference/event.md

@@ -21,9 +21,7 @@ A `'drain'` event is emitted whenever the `queueSize` reaches `0`.
 
 ## Event: `'needsDrain'`
 
-Similar to [`Piscina#needsDrain`](https://github.com/piscinajs/piscina#property-needsdrain-readonly);
-this event is triggered once the total capacity of the pool is exceeded
-by number of tasks enqueued that are pending of execution.
+Similar to [`Piscina#needsDrain`](https://github.com/piscinajs/piscina#property-needsdrain-readonly).
 
 ## Event: `'message'`
 

docs/docs/api-reference/property.md

@@ -68,11 +68,17 @@ The current number of tasks waiting to be assigned to a Worker thread.
 
 ## Property: `needsDrain` (readonly)
 
-Boolean value that specifies whether the capacity of the pool has
-been exceeded by the number of tasks submitted.
+Boolean value that specifies whether the pool requires draining before processing new tasks.
 
-This property is helpful to make decisions towards creating backpressure
-over the number of tasks submitted to the pool.
+This is `true` when there is not enough capacity to process more tasks
+without exceeding the maximum queue size or within the maximum number of threads.
+
+This is often calculated based on the number of threads (current and pending if number of threads is below `maxThreads` threshold),
+the maximum queue size, and the current number of in-flight tasks.
+
+This property is helpful to make decisions towards creating backpressure.
+
+> **Note**: Piscina does not buffer tasks beyond the `maxQueue` threshold, so this property is not a guarantee that the pool will accept more tasks.
 
 ## Property: `utilization` (readonly)
 

src/index.ts

@@ -452,7 +452,7 @@ class ThreadPool {
       this.histogram?.recordWaitTime(now - task.created)
       task.started = now;
       candidate[kWorkerData].postTask(task);
-      queueMicrotask(() => this._maybeDrain());
+      this._maybeDrain();
       // If candidate, let's try to distribute more tasks
       return true;
     }
@@ -515,7 +515,7 @@ class ThreadPool {
           resolve(result);
         }
 
-        queueMicrotask(this._maybeDrain.bind(this))
+        this._maybeDrain();
     });
 
     if (signal != null) {
@@ -558,7 +558,7 @@ class ThreadPool {
         this.taskQueue.push(taskInfo);
       }
 
-      queueMicrotask(this._maybeDrain.bind(this))
+      this._maybeDrain();
       return ret;
     }
 
@@ -578,7 +578,7 @@ class ThreadPool {
       }
     };
 
-    queueMicrotask(this._maybeDrain.bind(this))
+    this._maybeDrain();
     return ret;
   }
 
@@ -594,16 +594,26 @@ class ThreadPool {
      * since we want to avoid creating tasks that can't execute
      * immediately in order to provide back pressure to the task source.
      */
-    const { maxCapacity } = this;
+    const { maxCapacity, } = this;
     const currentUsage = this.workers.getCurrentUsage();
+    const maxQueueSize = this.options.maxQueue;
+    const queueSize = this.publicInterface.queueSize;
 
-    if (maxCapacity === currentUsage) {
+    if (this._needsDrain === true) {
+      if (queueSize === 0) {
+        this._needsDrain = false;
+        queueMicrotask(() => this.publicInterface.emit('drain'));
+      }
+      return;
+    }
+
+    // Attempting to provide a similar behaviour to a Writable stream
+    // if maxquesize is already reached, let's attempt to inform that the
+    // queue needs drain before handling more tasks
+    if (maxCapacity === currentUsage && queueSize === maxQueueSize) {
       this._needsDrain = true;
       queueMicrotask(() => this.publicInterface.emit('needsDrain'));
-    } else if (maxCapacity > currentUsage && this._needsDrain) {
-      this._needsDrain = false;
-      queueMicrotask(() => this.publicInterface.emit('drain'));
-    }
+    } 
   }
 
   async destroy () {
@@ -853,7 +863,7 @@ export default class Piscina<Exports extends Record<string, (payload: any) => an
 
   get queueSize () : number {
     const pool = this.#pool;
-    return Math.max(pool.taskQueue.size - pool.pendingCapacity(), 0);
+    return Math.max((pool.taskQueue.size + pool.skipQueue.length) - pool.pendingCapacity(), 0);
   }
 
   get completed () : number {

test/abort-task.test.ts

@@ -50,7 +50,7 @@ test('tasks can be aborted through EventEmitter before running', async () => {
   const ee = new EventEmitter();
   const task1 = pool.run(bufs[0]);
   const abortable = pool.run(bufs[1], { signal: ee });
-  assert.strictEqual(pool.queueSize, 0); // Means it's running
+  assert.strictEqual(pool.queueSize, 1); // Means it's running and abortable enqueued
   assert.rejects(abortable, /The task has been aborted/);
 
   ee.emit('abort');
@@ -75,7 +75,7 @@ test('abortable tasks will not share workers (abortable posted second)', async (
   const task1 = pool.run(bufs[0]);
   const ee = new EventEmitter();
   assert.rejects(pool.run(bufs[1], { signal: ee }), /The task has been aborted/);
-  assert.strictEqual(pool.queueSize, 0);
+  assert.strictEqual(pool.queueSize, 1);
 
   ee.emit('abort');
 

test/post-task.test.ts

@@ -2,6 +2,7 @@ import assert from 'node:assert';
 import { test, TestContext } from 'node:test';
 import { MessageChannel } from 'node:worker_threads';
 import { resolve } from 'node:path';
+
 import Piscina from '..';
 import { getAvailableParallelism } from '../dist/common';
 
@@ -77,36 +78,22 @@ test('postTask() validates abortSignal', () => {
     /signal argument must be an object/);
 });
 
-test('Piscina emits drain', async (t: TestContext) => {
-  const pool = new Piscina({
-    filename: resolve(__dirname, 'fixtures/eval.js'),
-    maxThreads: 1
-  });
-
-  t.plan(2);
-
-  pool.on('drain', () => {
-    t.assert.ok(true);
-    t.assert.ok(!pool.needsDrain);
-  });
-
-  await Promise.all([pool.run('123'), pool.run('123'), pool.run('123')]);
-});
-
-test('Piscina exposes/emits needsDrain to true when capacity is exceeded', async (t: TestContext) => {
+test('Piscina exposes/emits drain/needsDrain to true when capacity is exceeded', async (t: TestContext) => {
   const pool = new Piscina({
     filename: resolve(__dirname, 'fixtures/eval.js'),
     maxQueue: 3,
     maxThreads: 1
   });
 
-  t.plan(2);
+  t.plan(4);
 
   pool.once('drain', () => {
     t.assert.ok(true);
+    t.assert.ok(pool.queueSize < 3);
   });
   pool.once('needsDrain', () => {
     t.assert.ok(pool.needsDrain);
+    t.assert.equal(pool.queueSize, 3);
   });
 
   await Promise.all([