diff --git a/.changeset/stream-worker-chunks.md b/.changeset/stream-worker-chunks.md new file mode 100644 index 000000000..ce22833bf --- /dev/null +++ b/.changeset/stream-worker-chunks.md @@ -0,0 +1,5 @@ +--- +'@doc-kit/core': minor +--- + +`worker.stream` accepts an async iterable of batches as well as an array, starting a batch's work while earlier ones still run and reading the next batch once there is room for it. 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. diff --git a/packages/core/src/threading/__tests__/index.test.mjs b/packages/core/src/threading/__tests__/index.test.mjs index 203278840..127c0edc7 100644 --- a/packages/core/src/threading/__tests__/index.test.mjs +++ b/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(); + } + }); }); diff --git a/packages/core/src/threading/__tests__/parallel.test.mjs b/packages/core/src/threading/__tests__/parallel.test.mjs index 4a2c3653b..4c87e871a 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,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} toInput + * @returns {Promise} + */ +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, []); + }); + }); }); diff --git a/packages/core/src/threading/constants.mjs b/packages/core/src/threading/constants.mjs index 9872c6686..1ea409ec2 100644 --- a/packages/core/src/threading/constants.mjs +++ b/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; diff --git a/packages/core/src/threading/index.mjs b/packages/core/src/threading/index.mjs index 0a2ed5375..7cbf5b2b3 100644 --- a/packages/core/src/threading/index.mjs +++ b/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() }, }); } diff --git a/packages/core/src/threading/parallel.mjs b/packages/core/src/threading/parallel.mjs index 136bac69f..3563cadc9 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'); @@ -71,7 +74,8 @@ export default function createParallelWorker( pool, configuration ) { - const { threads, chunkSize } = configuration; + const { threads: defaultThreads = 1, chunkSize } = configuration; + const threads = Math.max(Number(defaultThreads), 1); const { name } = generator; @@ -79,54 +83,164 @@ export default function createParallelWorker( /** * Processes items in parallel, yielding results in submission order. * + * The items can also arrive in batches, as a dependency streams them: 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]; + + // The chunks the batches read so far were split into + let distributed = 0; + + /** + * Splits a batch into chunks. + * + * @param {T[]} batch - Items to split + * @returns {number[][]} Each chunk's indices into the batch + */ + 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)) { + const result = await generator.processChunk(batch, indices, extra); + + yield result; + } } return; } - // Start every chunk immediately, but emit results in submission order. - // Handle errors as each task settles so a later failure cannot 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) - ); + // Start only a few chunks per thread at a time, reading the next batch + // once there is room, so that generators running at once share the pool + // instead of one's chunks waiting behind all of another's. Emit results in + // submission order. Handle errors as each task settles so a later failure + // cannot become an unhandled rejection while earlier chunks are still + // running. + const tasks = []; + const limit = threads * CHUNKS_PER_THREAD; + + let running = 0; + let doneReading = false; + + // Resolved when a running chunk settles, while the reader waits for room + let room; + + // Resolved when a chunk starts or the reading ends, while the consumer + // waits for one + let added; + + /** + * Starts a batch's chunks, waiting for room whenever the limit is + * reached: an array's first chunks start right away. + * + * @param {T[]} batch - Items to start + */ + 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(); + } + }; + + /** + * Starts every batch's chunks, reading the next batch once there is room + * for them. + */ + const readBatches = async () => { + try { + if (!streamed) { + await enqueue(items); + + return; + } + + 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 failure reading the batches is thrown below, once every chunk + // submitted before it has been 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; + // Emitted chunks are the consumer's to keep or let go + 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, });