feat: Queued task prefers ready worker (#1058)

Author: isaacparker0

src/index.ts

@@ -404,7 +404,8 @@ class ThreadPool {
                        this.taskQueue.shift() as TaskInfo;
 
       if (workers == null) {
-        workers = [...this.workers].map(workerInfo => workerInfo.interface);
+        workers = [...this.workers.readyItems, ...this.workers.pendingItems]
+          .map(workerInfo => workerInfo.interface);
       }
 
       const distributed = this._distributeTask(taskInfo, workers);

test/fixtures/wait-until-released.js

@@ -0,0 +1,31 @@
+'use strict';
+
+// A worker whose readiness and task execution the main thread can release on
+// demand, letting a test deterministically hold a worker in the pending state
+// or keep one busy.
+//
+// The worker shares an Int32Array with the main thread (passed as workerData).
+// Each slot is a gate: the worker blocks on Atomics.wait until the main thread
+// stores a non-zero value into that slot and notifies.
+//
+// The readiness gate is awaited at module top level. A worker is only marked
+// ready once its module finishes loading, so blocking here keeps it pending
+// until the main thread lets it through. The task gate then holds a running
+// task until released. The returned threadId lets a test see which worker ran
+// a task.
+
+const { threadId, workerData } = require('node:worker_threads');
+
+const READY_GATE = 0;
+const TASK_GATE = 1;
+
+const gates = new Int32Array(workerData);
+
+Atomics.wait(gates, READY_GATE, 0);
+
+module.exports = ({ hold = false } = {}) => {
+  if (hold) {
+    Atomics.wait(gates, TASK_GATE, 0);
+  }
+  return threadId;
+};

test/queued-task-prefers-ready-worker.test.ts

@@ -0,0 +1,57 @@
+import * as assert from 'node:assert/strict';
+import { test } from 'node:test';
+import { resolve } from 'node:path';
+import Piscina from '..';
+
+const READY_SIGNAL = 0;
+const TASK_SIGNAL = 1;
+
+test('a queued task prefers a ready worker over a still-pending one', async () => {
+  const gates = new Int32Array(
+    new SharedArrayBuffer(2 * Int32Array.BYTES_PER_ELEMENT)
+  );
+
+  Atomics.store(gates, READY_SIGNAL, 1);
+
+  const pool = new Piscina({
+    filename: resolve(__dirname, 'fixtures/wait-until-released.js'),
+    minThreads: 1,
+    maxThreads: 2,
+    concurrentTasksPerWorker: 1,
+    workerData: gates.buffer
+  });
+
+  // One ready, idle worker.
+  const warmThreadId = await pool.run({});
+
+  // Keep the next worker from becoming ready.
+  Atomics.store(gates, READY_SIGNAL, 0);
+
+  // Occupy the ready worker, then queue a task behind it. A second worker
+  // spawns but stays pending.
+  const held = pool.run({ hold: true });
+  const queued = pool.run({});
+
+  // Free the ready worker while the second is still pending: the queued task
+  // should land on the ready worker.
+  Atomics.store(gates, TASK_SIGNAL, 1);
+  Atomics.notify(gates, TASK_SIGNAL, Infinity);
+
+  const heldThreadId = await held;
+
+  // Let the second worker finish so the pool can shut down (and the queued task
+  // resolves even if it was wrongly routed there).
+  Atomics.store(gates, READY_SIGNAL, 1);
+  Atomics.notify(gates, READY_SIGNAL, Infinity);
+
+  const queuedThreadId = await queued;
+
+  assert.strictEqual(heldThreadId, warmThreadId);
+  assert.strictEqual(
+    queuedThreadId,
+    warmThreadId,
+    'queued task should go to the ready worker, not the still-pending one'
+  );
+
+  await pool.destroy();
+});