diff --git a/.changeset/stream-worker-chunks.md b/.changeset/stream-worker-chunks.md new file mode 100644 index 000000000..34312584f --- /dev/null +++ b/.changeset/stream-worker-chunks.md @@ -0,0 +1,5 @@ +--- +'@doc-kit/core': patch +--- + +`worker.stream` also accepts an async iterable of batches, starting each batch's chunks while earlier ones still run. It keeps only a few chunks per thread in flight so concurrent generators share the worker pool, and still yields results in submission order. diff --git a/packages/core/src/threading/__tests__/parallel.test.mjs b/packages/core/src/threading/__tests__/parallel.test.mjs index 4a2c3653b..0f6cf1200 100644 --- a/packages/core/src/threading/__tests__/parallel.test.mjs +++ b/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,43 @@ async function collectChunks(generator) { return chunks; } +/** + * Streams 3x as many one-item chunks as two threads may hold in flight and + * checks the in-flight count peaks at exactly the limit, with results in order. + * + * @param {(items: number[]) => number[] | AsyncIterable} toInput + * @returns {Promise} + */ +async function checkChunksPerThread(toInput) { + const limit = 2 * CHUNKS_PER_THREAD; + const items = Array.from({ length: limit * 3 }, (_, index) => index); + + let running = 0; + let mostRunning = 0; + + const pool = { + run: async task => { + running++; + mostRunning = Math.max(mostRunning, running); + + await new Promise(resolve => setImmediate(resolve)); + + running--; + + return task.input; + }, + }; + const worker = createParallelWorker('ast-js', astJsGenerator, pool, { + threads: 2, + chunkSize: 1, + }); + + const results = await collectStream(worker.stream(toInput(items), {})); + + strictEqual(mostRunning, limit); + deepStrictEqual(results, items); +} + describe('createParallelWorker', () => { it('should create a ParallelWorker with stream method', async () => { const pool = createWorkerPool(2); @@ -230,6 +268,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 +301,162 @@ 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', + { 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(); + + 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', + { 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, []); + }); + }); }); diff --git a/packages/core/src/threading/constants.mjs b/packages/core/src/threading/constants.mjs index 9872c6686..21f5178a2 100644 --- a/packages/core/src/threading/constants.mjs +++ b/packages/core/src/threading/constants.mjs @@ -22,3 +22,9 @@ 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; + +/** + * Max chunks of one generator in flight per thread: enough to keep threads + * busy, few enough that concurrent generators share the pool. + */ +export const CHUNKS_PER_THREAD = 3; diff --git a/packages/core/src/threading/index.mjs b/packages/core/src/threading/index.mjs index 0a2ed5375..e3590545a 100644 --- a/packages/core/src/threading/index.mjs +++ b/packages/core/src/threading/index.mjs @@ -32,6 +32,8 @@ export default function createWorkerPool(threads) { maxYoungGenerationSizeMb: WORKER_GEN_SIZE, }, recordTiming: false, + // Start tasks in submission order, since results are yielded in that order + stricterFIFO: true, workerData: { logLevel: logger.getLogLevel() }, }); } diff --git a/packages/core/src/threading/parallel.mjs b/packages/core/src/threading/parallel.mjs index 136bac69f..494264369 100644 --- a/packages/core/src/threading/parallel.mjs +++ b/packages/core/src/threading/parallel.mjs @@ -1,6 +1,9 @@ 'use strict'; import logger from '#logger/index.mjs'; +import { isAsyncIterable } from '#utils/misc.mjs'; + +import { CHUNKS_PER_THREAD } from './constants.mjs'; const parallelLogger = logger.child('parallel'); @@ -79,54 +82,155 @@ export default function createParallelWorker( /** * Processes items in parallel, yielding results in submission order. * + * Items can also arrive as an async iterable of batches: a batch's chunks + * start while earlier ones still run, and the next batch is read once + * there is room for its chunks. + * * @template T - * @param {T[]} items - Items to process + * @param {T[] | AsyncIterable} items - Items to process * @param {object} extra - Extra options * @yields {R[]} Chunk results in input order */ async *stream(items, extra) { - if (items.length === 0) { + const streamed = isAsyncIterable(items); + + if (!streamed && items.length === 0) { return; } - const chunks = createChunks(items.length, chunkSize); + const batches = streamed ? items : [items]; + + // Total chunks across the batches read so far, for logging + let distributed = 0; + + /** + * Splits a batch into chunks of indices. + * + * @param {T[]} batch + * @returns {number[][]} + */ + const distribute = batch => { + const chunks = createChunks(batch.length, chunkSize); - parallelLogger.debug( - `Distributing ${items.length} items across ${chunks.length} chunks`, - { generator: name, chunks: chunks.length, chunkSize, threads } - ); + distributed += chunks.length; + + parallelLogger.debug( + `Distributing ${batch.length} items across ${chunks.length} chunks`, + { generator: name, chunks: chunks.length, chunkSize, threads } + ); + + return chunks; + }; // On this thread, one chunk at a time, so their work isn't all held at once - if (threads <= 1 || items.length <= 2) { - for (const indices of chunks) { - yield await generator.processChunk(items, indices, extra); + if (threads <= 1 || (!streamed && items.length <= 2)) { + for await (const batch of batches) { + for (const indices of distribute(batch)) { + yield await generator.processChunk(batch, indices, extra); + } } return; } - // Start every chunk immediately, but emit results in submission order. - // Handle errors as each task settles so a later failure cannot become an + // Keep only a few chunks per thread in flight so generators running at + // once share the pool, and emit results in submission order. Errors are + // captured as each task settles so a later failure can't become an // unhandled rejection while earlier chunks are still running. - const tasks = chunks.map(indices => { - const task = pool.run( - createTask(items, indices, extra, configuration, specifier, name) - ); + const tasks = []; + const limit = threads * CHUNKS_PER_THREAD; + + let running = 0; + let doneReading = false; + + // Resolved when a chunk settles (the reader waits on it for room) + let room; + + // Resolved when a chunk starts or reading ends (the consumer waits on it) + let added; + + /** + * Starts a batch's chunks, waiting for room whenever the limit is hit. + * + * @param {T[]} batch + */ + const enqueue = async batch => { + for (const indices of distribute(batch)) { + while (running >= limit) { + room = Promise.withResolvers(); + + await room.promise; + } + + running++; + + const task = pool.run( + createTask(batch, indices, extra, configuration, specifier, name) + ); + + tasks.push( + task + .then( + result => ({ result }), + error => ({ error }) + ) + .finally(() => { + running--; + room?.resolve(); + }) + ); + + added?.resolve(); + } + }; + + /** + * Reads every batch and starts its chunks. + */ + const readBatches = async () => { + try { + // An array's first chunks start synchronously + if (!streamed) { + return await enqueue(items); + } + + for await (const batch of items) { + await enqueue(batch); + } + } finally { + doneReading = true; + added?.resolve(); + } + }; - return task.then( - result => ({ result }), - error => ({ error }) - ); - }); + const reading = readBatches(); + + // A read failure is rethrown below, after every chunk before it is emitted + reading.catch(() => {}); + + for (let index = 0; ; index++) { + while (index >= tasks.length) { + if (doneReading) { + await reading; + + return; + } + + added = Promise.withResolvers(); + + await added.promise; + } + + const outcome = await tasks[index]; - for (const [index, task] of tasks.entries()) { - const outcome = await task; + // Drop our reference so emitted results can be collected + tasks[index] = undefined; if ('error' in outcome) { throw outcome.error; } - parallelLogger.debug(`Chunk ${index + 1}/${chunks.length} completed`, { + parallelLogger.debug(`Chunk ${index + 1}/${distributed} completed`, { generator: name, });