perf: let worker.stream take chunks as they arrive

.changeset/stream-worker-chunks.md

@@ -0,0 +1,5 @@
+---
+'@doc-kit/core': minor
+---
+
+`worker.stream` accepts an async iterable of batches, starting each one's work as it arrives, as well as an array. It keeps only a few chunks per thread in flight, so generators running at once share the worker pool, starts them in the order they were submitted and yields their results in that order. `toBatches`, from `@doc-kit/core/utils/generators.mjs`, gives an input that is either whole or streamed as batches.

packages/core/src/threading/__tests__/index.test.mjs

@@ -128,4 +128,14 @@ describe('createWorkerPool', () => {
       await pool.destroy();
     }
   });
+
+  it('should start tasks in the order they were submitted', async () => {
+    const pool = createWorkerPool(1);
+
+    try {
+      strictEqual(pool.options.stricterFIFO, true);
+    } finally {
+      await pool.destroy();
+    }
+  });
 });

packages/core/src/threading/__tests__/parallel.test.mjs

@@ -2,6 +2,7 @@ import { deepStrictEqual, ok, rejects, strictEqual } from 'node:assert';
 import { describe, it } from 'node:test';
 
 import { loadGenerator } from '../../generators/loader.mjs';
+import { CHUNKS_PER_THREAD } from '../constants.mjs';
 import createWorkerPool from '../index.mjs';
 import createParallelWorker from '../parallel.mjs';
 
@@ -42,6 +43,52 @@ async function collectChunks(generator) {
   return chunks;
 }
 
+/**
+ * Streams three times as many one-item chunks as two threads can run at once
+ * through a pool that holds each task until it is settled, newest first, and
+ * checks before each that as many run as there is room for: no more, and no
+ * fewer while any are left to start. Settling the newest first holds back the
+ * oldest chunk, so no result can be emitted until the last: the room must come
+ * from the settling alone. The results still come out in submission order.
+ * A stream that lost count of the held chunks would never end, so its tests
+ * need a timeout.
+ *
+ * @param {(items: number[]) => number[] | AsyncIterable<number[]>} toInput
+ * @returns {Promise<void>}
+ */
+async function checkChunksPerThread(toInput) {
+  const limit = 2 * CHUNKS_PER_THREAD;
+  const items = Array.from({ length: limit * 3 }, (_, index) => index);
+  const started = [];
+  const running = [];
+  const pool = {
+    run: task =>
+      new Promise(resolve => {
+        started.push(...task.input);
+        running.push(() => resolve(task.input));
+      }),
+  };
+  const worker = createParallelWorker('ast-js', astJsGenerator, pool, {
+    threads: 2,
+    chunkSize: 1,
+  });
+
+  const result = collectStream(worker.stream(toInput(items), {}));
+
+  for (let settled = 0; settled < items.length; settled++) {
+    await new Promise(resolve => setImmediate(resolve));
+
+    strictEqual(running.length, Math.min(limit, items.length - settled));
+
+    running.pop()();
+  }
+
+  const results = await result;
+
+  deepStrictEqual(started, items);
+  deepStrictEqual(results, items);
+}
+
 describe('createParallelWorker', () => {
   it('should create a ParallelWorker with stream method', async () => {
     const pool = createWorkerPool(2);
@@ -230,6 +277,12 @@ describe('createParallelWorker', () => {
     await rejects(result, /chunk failed/);
   });
 
+  it(
+    'should start a few chunks per thread at a time, and the next as each settles',
+    { timeout: 10_000 },
+    () => checkChunksPerThread(items => items)
+  );
+
   it('should process one chunk at a time, in order, with a single thread', async () => {
     let running = 0;
     let mostRunning = 0;
@@ -257,4 +310,167 @@ describe('createParallelWorker', () => {
     strictEqual(mostRunning, 1);
     deepStrictEqual(chunks, [[1, 2], [3, 4], [5]]);
   });
+
+  describe('with streamed batches', () => {
+    it(
+      'starts a batch before the next one arrives',
+      // The second batch only arrives once the first result is read, so a
+      // stream that waited for it would never give one
+      { timeout: 10_000 },
+      async () => {
+        const started = [];
+        const pool = {
+          run: async task => {
+            started.push(task.input);
+
+            return task.input;
+          },
+        };
+        const worker = createParallelWorker('ast-js', astJsGenerator, pool, {
+          threads: 2,
+          chunkSize: 2,
+        });
+        const gate = Promise.withResolvers();
+
+        async function* batches() {
+          yield [1, 2, 3];
+          await gate.promise;
+          yield [4];
+        }
+
+        const reader = worker.stream(batches(), {});
+        const first = await reader.next();
+
+        // The second batch hasn't arrived, but the first one's chunks ran
+        deepStrictEqual(first.value, [1, 2]);
+        deepStrictEqual(started, [[1, 2], [3]]);
+
+        gate.resolve();
+
+        const rest = await collectChunks(reader);
+
+        deepStrictEqual(rest, [[3], [4]]);
+      }
+    );
+
+    it(
+      'yields in submission order across batches',
+      // The test waits for every chunk to start, so it would wait forever on a
+      // stream that held one back
+      { timeout: 10_000 },
+      async () => {
+        const resolvers = [];
+        const pool = {
+          run: task =>
+            new Promise(resolve => {
+              resolvers.push(() => resolve(task.input));
+            }),
+        };
+        const worker = createParallelWorker('ast-js', astJsGenerator, pool, {
+          threads: 2,
+          chunkSize: 1,
+        });
+
+        async function* batches() {
+          yield [0, 1];
+          yield [2];
+        }
+
+        const result = collectStream(worker.stream(batches(), {}));
+
+        while (resolvers.length < 3) {
+          await new Promise(resolve => setImmediate(resolve));
+        }
+
+        for (const resolve of resolvers.toReversed()) {
+          resolve();
+        }
+
+        const results = await result;
+
+        deepStrictEqual(results, [0, 1, 2]);
+      }
+    );
+
+    it('throws a failure reading the batches after the chunks before it', async () => {
+      const pool = { run: async task => task.input };
+      const worker = createParallelWorker('ast-js', astJsGenerator, pool, {
+        threads: 2,
+        chunkSize: 1,
+      });
+
+      async function* batches() {
+        yield [1, 2];
+        throw new Error('source failed');
+      }
+
+      const chunks = [];
+
+      await rejects(async () => {
+        for await (const chunk of worker.stream(batches(), {})) {
+          chunks.push(chunk);
+        }
+      }, /source failed/);
+
+      deepStrictEqual(chunks, [[1], [2]]);
+    });
+
+    it(
+      'starts a few chunks per thread at a time, across batches',
+      { timeout: 10_000 },
+      () =>
+        checkChunksPerThread(async function* (items) {
+          for (let index = 0; index < items.length; index += 5) {
+            yield items.slice(index, index + 5);
+          }
+        })
+    );
+
+    it('runs each batch on this thread as it arrives with a single thread', async () => {
+      const events = [];
+      const generator = {
+        name: 'sequential',
+        async processChunk(items, indices) {
+          events.push(`chunk ${indices.map(index => items[index])}`);
+
+          return indices.map(index => items[index]);
+        },
+      };
+      const worker = createParallelWorker('sequential', generator, null, {
+        threads: 1,
+        chunkSize: 2,
+      });
+
+      async function* batches() {
+        events.push('batch 1');
+        yield [1, 2, 3];
+        events.push('batch 2');
+        yield [4];
+      }
+
+      const chunks = await collectChunks(worker.stream(batches(), {}));
+
+      deepStrictEqual(chunks, [[1, 2], [3], [4]]);
+      deepStrictEqual(events, [
+        'batch 1',
+        'chunk 1,2',
+        'chunk 3',
+        'batch 2',
+        'chunk 4',
+      ]);
+    });
+
+    it('yields nothing for no batches', async () => {
+      const worker = createParallelWorker('ast-js', astJsGenerator, null, {
+        threads: 2,
+        chunkSize: 1,
+      });
+
+      async function* batches() {}
+
+      const chunks = await collectChunks(worker.stream(batches(), {}));
+
+      deepStrictEqual(chunks, []);
+    });
+  });
 });

packages/core/src/threading/constants.mjs

@@ -22,3 +22,10 @@ export const WORKER_HEAP_SIZE = 512;
  * over half of a worker's heap. Node.js's `--max-semi-space-size` replaces it.
  */
 export const WORKER_GEN_SIZE = 96;
+
+/**
+ * How many chunks of one generator can be started, running or queued, per
+ * thread at a time: enough to keep every thread busy while more are submitted,
+ * few enough that generators running at once share the pool.
+ */
+export const CHUNKS_PER_THREAD = 3;

packages/core/src/threading/index.mjs

@@ -32,6 +32,10 @@ export default function createWorkerPool(threads) {
       maxYoungGenerationSizeMb: WORKER_GEN_SIZE,
     },
     recordTiming: false,
+    // A task taken from the queue that no worker can run yet goes back to the
+    // front, not the back, so tasks start in the order they were submitted
+    // and results, yielded in that order, are not held up
+    stricterFIFO: true,
     workerData: { logLevel: logger.getLogLevel() },
   });
 }