From a5992f743c8f50d0e572f3056d514b81ee59a9f5 Mon Sep 17 00:00:00 2001 From: Claudio Wunder Date: Tue, 3 Nov 2026 19:02:17 +0100 Subject: [PATCH 1/2] perf(core): stream a generator's output to the generators that read it Every dependency edge was a full barrier: the cache collected all of a generator's chunks before its dependents started, so each stage of a build waited for the whole of the one before it. A generator can now set streamsInput: true to receive its dependency's output as an async iterable of chunks, read as they are produced. With one thread nothing runs alongside, so it still gets one array, and a dependency that doesn't stream arrives as a single chunk. Nothing sets it yet. The cache runs a streamed generator once, recording its chunks so any number of readers each read all of them, from the first, while it still runs. Once every reader has started, it lets go of a chunk as soon as all of them have read it, rather than holding the output until the last one is done. Assisted-by: Claude Opus 5.5 Assisted-by: Claude Sonnet 5.5 --- .changeset/stream-input.md | 5 + docs/creating-generators.md | 54 ++- packages/core/src/__tests__/caching.test.mjs | 416 +++++++++++++++++- .../core/src/__tests__/generators.test.mjs | 132 ++++++ packages/core/src/caching.mjs | 134 ++++-- packages/core/src/generators.mjs | 22 +- packages/core/src/generators/types.d.ts | 25 +- .../core/src/utils/__tests__/replay.test.mjs | 86 ++++ packages/core/src/utils/replay.mjs | 152 +++++++ 9 files changed, 985 insertions(+), 41 deletions(-) create mode 100644 .changeset/stream-input.md create mode 100644 packages/core/src/utils/__tests__/replay.test.mjs create mode 100644 packages/core/src/utils/replay.mjs diff --git a/.changeset/stream-input.md b/.changeset/stream-input.md new file mode 100644 index 000000000..977c18f1a --- /dev/null +++ b/.changeset/stream-input.md @@ -0,0 +1,5 @@ +--- +'@doc-kit/core': minor +--- + +Generators can start on their dependency's first chunk instead of waiting for all of it: with `streamsInput: true`, a generator receives its dependency's output as an async iterable of chunks (with one thread, still whole), which `worker.stream` accepts. The cache records a streamed generator once for all the generators reading it, each from the first chunk, and lets go of a chunk once every one of them has read it. diff --git a/docs/creating-generators.md b/docs/creating-generators.md index c2912d335..62aa810d3 100644 --- a/docs/creating-generators.md +++ b/docs/creating-generators.md @@ -330,9 +330,61 @@ export async function* generate(input, worker) { ### Benefits of Streaming - **Reduced memory usage** - Process data in chunks -- **Earlier downstream starts** - Next generator can begin before this one finishes +- **Earlier downstream starts** - A generator that streams its input (see below) + begins on the first chunk, before this one finishes - **Better parallelism** - Multiple generators can work simultaneously +### Streaming Input + +By default, a generator receives its dependency's output once it is complete, +collected into one array. With `streamsInput: true`, it receives it as it is +produced instead: `input` is an async iterable of the dependency's chunks, and +if the dependency streams its results too, the generator starts as soon as the +dependency does. + +```mjs displayName="index.mjs" +import { generate, processChunk } from './generate.mjs'; + +export default { + name: 'per-chunk-generator', + + description: 'Processes each chunk of metadata as it arrives', + + dependsOn: '@doc-kit/core/metadata', + + hasParallelProcessor: true, + + streamsInput: true, + + generate, + processChunk, +}; +``` + +`worker.stream` takes the chunks as they arrive, starting each one's work while +the next ones are still being produced: + +```mjs displayName="generate.mjs" +export async function* generate(input, worker) { + for await (const chunkResult of worker.stream(input, {})) { + yield chunkResult; + } +} +``` + +- A dependency that does not stream arrives as a single chunk, once it is + complete. +- With one thread, nothing would run alongside, so the input arrives whole, as + it does without `streamsInput`: accept both, as `worker.stream` does. +- With more than one thread, `input` can be iterated only once, and passing it + to `worker.stream` iterates it. If you need it twice, collect it into an + array first. +- The chunks are the dependency's own, in the order it yields them. +- Work that needs the whole input, such as sorting or writing one file, has to + wait for the last chunk. +- Other generators reading the same dependency are not affected: each gets + every chunk, from the first. + ### Non-Streaming Generators Some generators must collect all input before processing. diff --git a/packages/core/src/__tests__/caching.test.mjs b/packages/core/src/__tests__/caching.test.mjs index 5dc306e84..223062936 100644 --- a/packages/core/src/__tests__/caching.test.mjs +++ b/packages/core/src/__tests__/caching.test.mjs @@ -1,8 +1,121 @@ -import { deepStrictEqual, strictEqual } from 'node:assert'; +import { deepStrictEqual, rejects, strictEqual, throws } from 'node:assert'; import { describe, it } from 'node:test'; +import { setFlagsFromString } from 'node:v8'; +import { runInNewContext } from 'node:vm'; import { collectAsyncGenerator, createCache } from '../caching.mjs'; +setFlagsFromString('--expose-gc'); + +// Collects garbage, to tell what is let go +const gc = runInNewContext('gc'); + +/** + * Reads every chunk of an async iterable. + * + * @template T + * @param {AsyncIterable} iterable + * @returns {Promise} + */ +async function collectChunks(iterable) { + const chunks = []; + + for await (const chunk of iterable) { + chunks.push(chunk); + } + + return chunks; +} + +/** + * Reads up to `count` chunks of items, keeping only the items' values, so that + * the test holds none of the chunks. + * + * @param {AsyncIterator>} reader + * @param {number} [count] + * @returns {Promise} + */ +async function readValues(reader, count = Infinity) { + const values = []; + + for (let i = 0; i < count; i++) { + const { value, done } = await reader.next(); + + if (done) { + break; + } + + values.push(...value.map(item => item.value)); + } + + return values; +} + +/** + * Makes chunks of one item each that tell when they are let go. + * + * @returns {{ released: Set, chunkOf: (value: number) => Array<{ value: number }> }} + */ +function trackReleases() { + const released = new Set(); + const registry = new FinalizationRegistry(value => released.add(value)); + + return { + released, + + /** + * A chunk holding an item that tells when it is let go. + * + * @param {number} value + * @returns {Array<{ value: number }>} + */ + chunkOf: value => { + const item = { value }; + + registry.register(item, value); + + return [item]; + }, + }; +} + +/** + * Collects garbage until `done` holds, at most ten times. + * + * @param {() => boolean} done + * @returns {Promise} + */ +async function collectGarbage(done) { + for (let i = 0; i < 10 && !done(); i++) { + await new Promise(resolve => setImmediate(resolve)); + gc(); + } +} + +/** + * Runs `fn`, and gives the rejections left unhandled while it ran. + * + * @param {() => Promise} fn + * @returns {Promise} + */ +async function unhandledRejectionsOf(fn) { + const unhandled = []; + const onUnhandled = error => unhandled.push(error); + + process.on('unhandledRejection', onUnhandled); + + try { + await fn(); + + // A rejection left unhandled is reported once the microtasks have run + await new Promise(resolve => setImmediate(resolve)); + } finally { + process.off('unhandledRejection', onUnhandled); + } + + return unhandled; +} + describe('caching', () => { describe('collectAsyncGenerator', () => { it('should collect all chunks into a flat array', async () => { @@ -106,5 +219,306 @@ describe('caching', () => { await cache.consume('a'); strictEqual(cache.has('a'), false); }); + + it( + 'should stream chunks before the generator completes', + // The generator waits for its first chunk to be read, so a stream that + // waited for the generator to complete would never give it + { timeout: 10_000 }, + async () => { + const cache = createCache(); + const gate = Promise.withResolvers(); + + async function* gen() { + yield [1]; + await gate.promise; + yield [2]; + } + + cache.store('a', gen()); + + const stream = await cache.stream('a'); + const reader = stream[Symbol.asyncIterator](); + const first = await reader.next(); + + // The first chunk arrives while the generator still waits on the gate + deepStrictEqual(first.value, [1]); + + gate.resolve(); + + const rest = await collectChunks(reader); + + deepStrictEqual(rest, [[2]]); + } + ); + + it('should run a generator once for streaming and collecting consumers', async () => { + const cache = createCache(); + + let iterations = 0; + + async function* gen() { + iterations++; + yield [1, 2]; + yield [3]; + } + + cache.store('a', gen()); + + const [streamed, collected] = await Promise.all([ + cache.stream('a').then(collectChunks), + cache.consume('a'), + ]); + + deepStrictEqual(streamed, [[1, 2], [3]]); + deepStrictEqual(collected, [1, 2, 3]); + strictEqual(iterations, 1); + }); + + it('should replay every chunk to a reader that starts late', async () => { + const cache = createCache(); + + async function* gen() { + yield [1]; + yield [2]; + } + + cache.store('a', gen()); + + await cache.consume('a'); + + const reader = await cache.stream('a'); + const chunks = await collectChunks(reader); + + deepStrictEqual(chunks, [[1], [2]]); + }); + + it('should stream a non-streaming result as a single chunk', async () => { + const cache = createCache(); + + cache.store('a', Promise.resolve([1, 2])); + + const reader = await cache.stream('a'); + const chunks = await collectChunks(reader); + + deepStrictEqual(chunks, [[1, 2]]); + }); + + it('should throw a failure to streaming readers after the chunks before it', async () => { + const cache = createCache(); + + async function* gen() { + yield [1]; + throw new Error('boom'); + } + + cache.store('a', gen()); + + const chunks = []; + + const unhandled = await unhandledRejectionsOf(() => + rejects(async () => { + const reader = await cache.stream('a'); + + for await (const chunk of reader) { + chunks.push(chunk); + } + }, /boom/) + ); + + deepStrictEqual(chunks, [[1]]); + deepStrictEqual(unhandled, []); + }); + + it('should handle a failure before the stream is read', async () => { + const cache = createCache(); + + const unhandled = await unhandledRejectionsOf(async () => { + cache.store('a', Promise.reject(new Error('boom'))); + + // The reader is asked for as a generator that streams its input + // starts, before it reads anything + await rejects(cache.stream('a'), /boom/); + }); + + deepStrictEqual(unhandled, []); + }); + + it('should start recording a streamed result before it is read', async () => { + const cache = createCache(); + + let started = false; + + async function* gen() { + started = true; + yield [1]; + } + + cache.store('a', gen()); + + const reader = await cache.stream('a'); + + await new Promise(resolve => setImmediate(resolve)); + + strictEqual(started, true); + + const chunks = await collectChunks(reader); + + deepStrictEqual(chunks, [[1]]); + }); + + it('should read a stream only once', async () => { + const cache = createCache(); + + async function* gen() { + yield [1]; + } + + cache.store('a', gen()); + + const stream = await cache.stream('a'); + const chunks = await collectChunks(stream); + + deepStrictEqual(chunks, [[1]]); + throws(() => stream[Symbol.asyncIterator](), /can only be read once/); + }); + + it('should give a reader every chunk even if it starts reading after the result was evicted', async () => { + const cache = createCache(); + + cache.populateConsumerCounts(['a'], () => undefined); + + async function* gen() { + yield [1]; + yield [2]; + yield [3]; + } + + cache.store('a', gen()); + + // Its only read: the result is evicted, and its chunks go once read + const stream = await cache.stream('a'); + + strictEqual(cache.has('a'), false); + + // Every chunk is recorded before the reader reads any + await new Promise(resolve => setImmediate(resolve)); + + const chunks = await collectChunks(stream); + + deepStrictEqual(chunks, [[1], [2], [3]]); + }); + + it('should let go of the chunks a reader has read, after the generator ended', async () => { + const cache = createCache(); + const { released, chunkOf } = trackReleases(); + + cache.populateConsumerCounts(['a'], () => undefined); + + async function* gen() { + yield chunkOf(1); + yield chunkOf(2); + yield chunkOf(3); + } + + cache.store('a', gen()); + + // Its only read: no more readers can start + const stream = await cache.stream('a'); + + // The generator ends before any chunk is read + await new Promise(resolve => setImmediate(resolve)); + + const reader = stream[Symbol.asyncIterator](); + const read = await readValues(reader, 2); + + await collectGarbage(() => released.has(1)); + + strictEqual(released.has(1), true); + + const rest = await readValues(reader); + + // The last chunk goes once the reader is done, while it is still held + await collectGarbage(() => released.has(3)); + + const end = await reader.next(); + + deepStrictEqual([...read, ...rest], [1, 2, 3]); + strictEqual(end.done, true); + strictEqual(released.has(3), true); + }); + + it('should let go of a chunk only once every streaming reader has read it', async () => { + const cache = createCache(); + const { released, chunkOf } = trackReleases(); + + // `x`, `y` and `z` all read `a` + cache.populateConsumerCounts(['x', 'y', 'z'], name => + name === 'a' ? undefined : 'a' + ); + + async function* gen() { + yield chunkOf(1); + yield chunkOf(2); + yield chunkOf(3); + } + + cache.store('a', gen()); + + const streams = []; + + // The third is the last read: no more readers can start + for (let i = 0; i < 3; i++) { + const stream = await cache.stream('a'); + + streams.push(stream); + } + + // The reader that falls behind started between the other two, so neither + // the first reader nor the last one alone tells what all have read + const [first, middle, last] = streams.map(stream => + stream[Symbol.asyncIterator]() + ); + + const firstRead = await readValues(first, 2); + const lastRead = await readValues(last, 2); + + // `middle` has yet to read the first chunk + await collectGarbage(() => released.has(1)); + + strictEqual(released.has(1), false); + + const middleRead = await readValues(middle, 2); + + await collectGarbage(() => released.has(1)); + + strictEqual(released.has(1), true); + + const firstRest = await readValues(first); + const middleRest = await readValues(middle); + const lastRest = await readValues(last); + + deepStrictEqual([...firstRead, ...firstRest], [1, 2, 3]); + deepStrictEqual([...middleRead, ...middleRest], [1, 2, 3]); + deepStrictEqual([...lastRead, ...lastRest], [1, 2, 3]); + }); + + it('should count a stream as a read for eviction', async () => { + const cache = createCache(); + + cache.populateConsumerCounts(['a'], () => undefined); + + async function* gen() { + yield [1]; + } + + cache.store('a', gen()); + + const reader = await cache.stream('a'); + const chunks = await collectChunks(reader); + + deepStrictEqual(chunks, [[1]]); + strictEqual(cache.has('a'), false); + }); }); }); diff --git a/packages/core/src/__tests__/generators.test.mjs b/packages/core/src/__tests__/generators.test.mjs index d3be80da3..07fb99d32 100644 --- a/packages/core/src/__tests__/generators.test.mjs +++ b/packages/core/src/__tests__/generators.test.mjs @@ -16,6 +16,9 @@ const streamOf = chunk => yield chunk; })(); +// Lets `gated-source` wait for its streaming reader +const gate = { opened: Promise.withResolvers(), openedBy: undefined }; + // Synthetic generators keyed by specifier (any string works as a specifier) const syntheticGenerators = { // Root streaming generator (no dependency) @@ -99,6 +102,81 @@ const syntheticGenerators = { return { all: input }; }, }, + // A source that only yields its second chunk once its streaming reader has + // received the first, or after a second, recording which happened first + 'gated-source': { + name: 'gated-source', + hasParallelProcessor: true, + generate: async () => { + record('gated-source'); + + return (async function* () { + yield [1]; + + let timer; + + const timeout = new Promise(resolve => { + timer = setTimeout(resolve, 1000, 'timeout'); + }); + + gate.openedBy = await Promise.race([gate.opened.promise, timeout]); + + clearTimeout(timer); + + yield [2]; + })(); + }, + }, + 'stream-reader': { + name: 'stream-reader', + dependsOn: 'gated-source', + streamsInput: true, + async *generate(input) { + record('stream-reader'); + + for await (const chunk of input) { + gate.opened.resolve('reader'); + + yield [{ streamed: chunk }]; + } + }, + }, + // A source that fails before yielding anything, read only by a generator + // that streams its input and starts reading it a turn of the event loop + // later, once a failure left unhandled would have been reported + 'failing-source': { + name: 'failing-source', + hasParallelProcessor: true, + generate: async () => { + throw new Error('source failed'); + }, + }, + 'slow-reader': { + name: 'slow-reader', + dependsOn: 'failing-source', + streamsInput: true, + async *generate(input) { + await new Promise(resolve => setImmediate(resolve)); + + for await (const chunk of input) { + yield chunk; + } + }, + }, + 'shape-reader': { + name: 'shape-reader', + dependsOn: 'metadata', + streamsInput: true, + generate: async input => ({ array: Array.isArray(input) }), + }, + 'collect-reader': { + name: 'collect-reader', + dependsOn: 'gated-source', + generate: async input => { + record('collect-reader'); + return { collected: input }; + }, + }, }; mock.module('../generators/loader.mjs', { @@ -245,4 +323,58 @@ describe('createGenerator orchestration', () => { { name: 'gen-a', markdown, ran: ['ast', 'metadata'] }, ]); }); + + it('streams a dependency to a generator that streams its input', async () => { + const { runGenerators } = createGenerator(); + + const results = await runGenerators({ + target: ['stream-reader', 'collect-reader'], + threads: 2, + }); + + // The reader got the first chunk while the source was still running + assert.equal(gate.openedBy, 'reader'); + assert.equal(runs['gated-source'], 1); + + assert.deepStrictEqual(results, [ + [{ streamed: [1] }, { streamed: [2] }], + { collected: [1, 2] }, + ]); + }); + + it('rejects with the failure of a dependency read as a stream', async () => { + const { runGenerators } = createGenerator(); + const unhandled = []; + const onUnhandled = error => unhandled.push(error); + + process.on('unhandledRejection', onUnhandled); + + try { + await assert.rejects( + runGenerators({ target: ['slow-reader'], threads: 2 }), + /source failed/ + ); + + // A rejection left unhandled is reported once the microtasks have run + await new Promise(resolve => setImmediate(resolve)); + } finally { + process.off('unhandledRejection', onUnhandled); + } + + assert.deepStrictEqual(unhandled, []); + }); + + it('gives a generator that streams its input the whole input with one thread', async () => { + const one = await createGenerator().runGenerators({ + target: ['shape-reader'], + threads: 1, + }); + const two = await createGenerator().runGenerators({ + target: ['shape-reader'], + threads: 2, + }); + + assert.deepStrictEqual(one, [{ array: true }]); + assert.deepStrictEqual(two, [{ array: false }]); + }); }); diff --git a/packages/core/src/caching.mjs b/packages/core/src/caching.mjs index 9ad72110e..bdfcf4dfd 100644 --- a/packages/core/src/caching.mjs +++ b/packages/core/src/caching.mjs @@ -2,6 +2,7 @@ import logger from './logger/index.mjs'; import { isAsyncIterable } from './utils/misc.mjs'; +import { createReplay } from './utils/replay.mjs'; const cachingLogger = logger.child('caching'); @@ -36,6 +37,17 @@ export const collectAsyncGenerator = async generator => { return results; }; +/** + * Gives a value as the single chunk of an async generator. + * + * @template T + * @param {T} value + * @returns {AsyncGenerator} + */ +async function* once(value) { + yield value; +} + /** * Creates the cache that backs the generator pipeline. It is the single owner * of every piece of caching the orchestrator needs. @@ -43,17 +55,23 @@ export const collectAsyncGenerator = async generator => { * Generator caching: each generator's pending result, so a generator and its * dependencies are only ever scheduled once. * - * Stream caching: collecting an async generator's chunks into a single array - * exactly once, shared by every consumer. + * Stream caching: running an async generator exactly once, recording its + * chunks for every consumer: those that collect it share a single array, and + * those that stream it read each chunk as soon as it is yielded. * * Dependency caching & eviction: tracking how many consumers still need each - * result and dropping it once the last one has read it, so intermediate results - * can be garbage collected while later generators keep running. + * result and dropping it once the last one has asked for it, and a streamed + * result's chunks once every reader has read them, so intermediate results can + * be garbage collected while later generators keep running. */ export const createCache = () => { /** @type {{ [key: string]: Promise | AsyncGenerator }} */ const results = {}; + // The replays of the results that stream, and those of them collected into + // one array, shared by the consumers that collect them + /** @type {Map>} */ + const replays = new Map(); /** @type {Map>} */ const collections = new Map(); @@ -66,27 +84,40 @@ export const createCache = () => { const remaining = {}; /** - * Collects an async generator's chunks, ensuring a single collection is - * shared across all consumers of the same key. + * Gets the replay of a streamed result, starting it the first time a + * consumer reads the result, so every consumer shares its one run. * - * @param {string} key - Cache key (resolved generator specifier) - * @param {AsyncGenerator} generator - Generator to collect - * @returns {Promise} Promise resolving to the collected items + * @param {string} name - Cache key (resolved generator specifier) + * @param {AsyncGenerator} generator - The result + * @returns {ReturnType} */ - const collect = (key, generator) => { - const hasKey = collections.has(key); + const replayOf = (name, generator) => { + if (!replays.has(name)) { + cachingLogger.debug(`Starting to record "${name}"`); - if (!hasKey) { - collections.set(key, collectAsyncGenerator(generator)); + replays.set(name, createReplay(generator)); } - cachingLogger.debug( - hasKey - ? `Using cached collection for "${key}"` - : `Starting collection for "${key}"` - ); + return replays.get(name); + }; + + /** + * Counts a read of a result, and evicts it once every expected consumer has + * read it. + * + * @param {string} name - Cache key (resolved generator specifier) + */ + const release = name => { + if (name in remaining && --remaining[name] <= 0) { + delete results[name]; + + // Every reader has started, so what they all read can go + replays.get(name)?.forget(); + replays.delete(name); + collections.delete(name); - return collections.get(key); + cachingLogger.debug(`Evicted "${name}" after final consumer`); + } }; return { @@ -165,19 +196,68 @@ export const createCache = () => { consume: async name => { const result = await results[name]; - const value = isAsyncIterable(result) - ? await collect(name, result) - : result; + let value = result; - // Evict once the last consumer has retrieved the result - if (name in remaining && --remaining[name] <= 0) { - delete results[name]; - collections.delete(name); + if (isAsyncIterable(result)) { + // Consumers that collect it share one array, read like a stream + if (!collections.has(name)) { + collections.set( + name, + collectAsyncGenerator(replayOf(name, result).read()) + ); + } - cachingLogger.debug(`Evicted "${name}" after final consumer`); + value = await collections.get(name); } + release(name); + return value; }, + + /** + * Gives a reader of a generator's result as it is produced: chunk by chunk + * if it streams, or as a single chunk once it resolves. As with `consume`, + * the result is awaited and its recording started right away, so the + * generator runs and its failure is handled before the reader starts + * reading, and this counts as one of its reads. + * + * @param {string} name - Cache key (resolved generator specifier) to stream + * @returns {Promise>} + */ + stream: async name => { + const result = await results[name]; + + let reader; + + if (isAsyncIterable(result)) { + reader = replayOf(name, result).read(); + } else { + // A result that doesn't stream is read as a single chunk + reader = once(result); + } + + release(name); + + let read = false; + + // Read once: what every reader has read is let go + return { + /** + * Reads every chunk, from the first, once. + * + * @returns {AsyncGenerator} + */ + [Symbol.asyncIterator]: () => { + if (read) { + throw new Error(`The output of "${name}" can only be read once`); + } + + read = true; + + return reader; + }, + }; + }, }; }; diff --git a/packages/core/src/generators.mjs b/packages/core/src/generators.mjs index b9959f4ae..b88809fb6 100644 --- a/packages/core/src/generators.mjs +++ b/packages/core/src/generators.mjs @@ -28,17 +28,20 @@ const createGenerator = () => { let pool; /** - * Gets the collected input from a dependency generator. + * Gets the input from a dependency generator: collected, or its chunks as the + * dependency yields them. * * @param {string | undefined} dependsOn - Resolved dependency specifier + * @param {boolean} [streamed] - Whether to give its chunks as they are yielded + * rather than collected * @returns {Promise} */ - const getDependencyInput = async dependsOn => { + const getDependencyInput = async (dependsOn, streamed) => { if (!dependsOn) { return undefined; } - return cache.consume(dependsOn); + return streamed ? cache.stream(dependsOn) : cache.consume(dependsOn); }; /** @@ -57,7 +60,7 @@ const createGenerator = () => { } const generator = generators.get(specifier); - const { name, generate, hasParallelProcessor } = generator; + const { name, generate, hasParallelProcessor, streamsInput } = generator; const dependsOn = inputOf.get(specifier); @@ -75,7 +78,12 @@ const createGenerator = () => { cache.store( specifier, (async () => { - const dependencyInput = await getDependencyInput(dependsOn); + // With one thread nothing runs alongside, so streaming would only keep + // every stage's data alive at once: each stage takes its input whole + const dependencyInput = await getDependencyInput( + dependsOn, + streamsInput && configuration.threads > 1 + ); generatorsLogger.debug(`Starting "${name}"`); @@ -91,8 +99,8 @@ const createGenerator = () => { // input while the generator runs, and the generator can let it go return generate(dependencyInput, await worker); })().then(result => { - // For streaming generators, "Completed" is logged when the cache - // finishes collecting, not here when the generator returns + // For streaming generators, the cache logs each chunk as it records + // it, not here when the generator returns if (!isAsyncIterable(result)) { generatorsLogger.debug(`Completed "${name}"`); } diff --git a/packages/core/src/generators/types.d.ts b/packages/core/src/generators/types.d.ts index 8bee0399e..51e7634cb 100644 --- a/packages/core/src/generators/types.d.ts +++ b/packages/core/src/generators/types.d.ts @@ -13,16 +13,18 @@ declare global { */ export interface ParallelWorker { /** - * Processes items in parallel across worker threads and yields results - * as each chunk completes. Enables downstream processing to begin + * Processes items in parallel across worker threads and yields each + * chunk's results in submission order, as soon as it and every chunk + * before it have completed. Enables downstream processing to begin * while upstream chunks are still being processed. * - * @param items - Items to process (determines chunk distribution) + * @param items - Items to process (determines chunk distribution), or + * batches of them as a streamed input yields them * @param opts - Additional options to pass to workers - * @yields Each chunk's results as they complete + * @yields Each chunk's results, in submission order */ stream( - items: T[], + items: T[] | AsyncIterable, opts?: Record ): AsyncGenerator; } @@ -70,6 +72,19 @@ declare global { hasParallelProcessor?: boolean; + /** + * Receive the dependency's output as it is produced, as an async iterable + * of its chunks, instead of collected into one array once it is complete, + * so this generator can start on the first chunk while the rest are still + * being produced. The async iterable can be iterated only once. A + * dependency that doesn't stream arrives as a single chunk, once it is + * complete. + * + * With one thread, the input still arrives collected into one array, as + * without `streamsInput`, so accept both. + */ + streamsInput?: boolean; + /** * The unified plugins this generator processes Markdown with, listed as * in the `markdown` option, with paths relative to its module. Each list diff --git a/packages/core/src/utils/__tests__/replay.test.mjs b/packages/core/src/utils/__tests__/replay.test.mjs new file mode 100644 index 000000000..2225e4beb --- /dev/null +++ b/packages/core/src/utils/__tests__/replay.test.mjs @@ -0,0 +1,86 @@ +import { deepStrictEqual, rejects } from 'node:assert'; +import { describe, it } from 'node:test'; + +import { createReplay } from '../replay.mjs'; + +/** + * Reads every chunk of an async iterable. + * + * @template T + * @param {AsyncIterable} iterable + * @returns {Promise} + */ +async function collectChunks(iterable) { + const chunks = []; + + for await (const chunk of iterable) { + chunks.push(chunk); + } + + return chunks; +} + +describe('createReplay', () => { + it('should give every reader every chunk, in order', async () => { + const gate = Promise.withResolvers(); + + async function* gen() { + yield [1]; + await gate.promise; + yield [2]; + yield [3]; + } + + const replay = createReplay(gen()); + + const early = collectChunks(replay.read()); + + await new Promise(resolve => setImmediate(resolve)); + gate.resolve(); + + const late = collectChunks(replay.read()); + + const earlyChunks = await early; + const lateChunks = await late; + + deepStrictEqual(earlyChunks, [[1], [2], [3]]); + deepStrictEqual(lateChunks, [[1], [2], [3]]); + }); + + it('should throw rather than give a reader a chunk that was let go', async () => { + async function* gen() { + yield [1]; + yield [2]; + } + + const replay = createReplay(gen()); + + await replay.ended; + + // With no reader left, every chunk goes + replay.forget(); + + await rejects(collectChunks(replay.read()), /Chunk 0 was let go/); + }); + + it('should run the generator to its end, and throw its failure to readers', async () => { + async function* gen() { + yield [1]; + + throw new Error('replayed failure'); + } + + const replay = createReplay(gen()); + const chunks = []; + + await rejects(replay.ended, /replayed failure/); + + await rejects(async () => { + for await (const chunk of replay.read()) { + chunks.push(chunk); + } + }, /replayed failure/); + + deepStrictEqual(chunks, [[1]]); + }); +}); diff --git a/packages/core/src/utils/replay.mjs b/packages/core/src/utils/replay.mjs new file mode 100644 index 000000000..52b555b31 --- /dev/null +++ b/packages/core/src/utils/replay.mjs @@ -0,0 +1,152 @@ +'use strict'; + +import { createSignal } from './misc.mjs'; + +/** + * Runs an async generator to its end, recording its chunks so that any number + * of readers can each read all of them, from the first, while it still runs. + * + * Once no more readers can start (see `forget`), a chunk every reader has read + * is let go. + * + * @template T + * @param {AsyncGenerator} generator - Generator to run + * @returns {{ ended: Promise, read: () => AsyncGenerator, forget: () => void }} + */ +export const createReplay = generator => { + /** @type {Array} */ + const chunks = []; + + // Where each reader is, by the index of the next chunk it reads + const readers = new Set(); + + let done = false; + let failure; + let forgotten = false; + let trimmed = 0; + + // Notified as each chunk is recorded, and at the end + const recorded = createSignal(); + + /** + * Lets go of the chunks every reader has read, once no more can start. + */ + const trim = () => { + if (!forgotten) { + return; + } + + let read = chunks.length; + + for (const reader of readers) { + read = Math.min(read, reader.index); + } + + for (; trimmed < read; trimmed++) { + chunks[trimmed] = undefined; + } + }; + + /** + * Runs the generator to its end, recording each chunk, whether or not its + * readers keep up. + */ + const record = async () => { + try { + for await (const chunk of generator) { + chunks.push(chunk); + recorded.notify(); + trim(); + } + } catch (error) { + failure = { error }; + + throw error; + } finally { + done = true; + recorded.notify(); + + // Readers share this closure: let the finished generator, and all it + // holds, go + generator = undefined; + } + }; + + /** + * Reads every chunk from where a reader is, those recorded and those still + * to come, then throws if the generator failed. + * + * @param {{ index: number }} reader + * @returns {AsyncGenerator} + */ + async function* readFrom(reader) { + try { + for (; ; reader.index++) { + while (reader.index >= chunks.length) { + if (failure) { + throw failure.error; + } + + if (done) { + return; + } + + await recorded.wait(); + } + + // A reader started after `forget` can be behind what was let go: fail + // rather than give it `undefined` for a chunk + if (reader.index < trimmed) { + throw new Error( + `Chunk ${reader.index} was let go before this reader read it` + ); + } + + const chunk = chunks[reader.index]; + + trim(); + + yield chunk; + } + } finally { + readers.delete(reader); + trim(); + } + } + + const ended = record(); + + // Its readers get a failure as they read; this keeps one with none from + // being reported as unhandled + ended.catch(() => {}); + + return { + /** + * Resolves once the generator has ended, or rejects with its failure. + */ + ended, + + /** + * Reads every chunk, from the first. The reader counts from when this is + * called, not from when it is first read, so nothing it hasn't read yet + * goes in the meantime. + * + * @returns {AsyncGenerator} + */ + read() { + const reader = { index: 0 }; + + readers.add(reader); + + return readFrom(reader); + }, + + /** + * Says no more readers will start, so chunks every reader has read can go. + */ + forget() { + forgotten = true; + trim(); + }, + }; +}; From 557f3193b6f04512b29ea72b9921503fa606f3d5 Mon Sep 17 00:00:00 2001 From: Claudio Wunder Date: Wed, 4 Nov 2026 19:07:09 +0100 Subject: [PATCH 2/2] refactor: add createSignal with the replay, its first caller worker.stream no longer waits on a signal, so createSignal comes with the replay, which wakes its readers with it as chunks are recorded. Assisted-by: Claude Opus 5.5 --- packages/core/src/utils/misc.mjs | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/packages/core/src/utils/misc.mjs b/packages/core/src/utils/misc.mjs index 92d91e01a..6b73c1960 100644 --- a/packages/core/src/utils/misc.mjs +++ b/packages/core/src/utils/misc.mjs @@ -70,3 +70,29 @@ export const deepMerge = (...objects) => { return result; }, {}); }; + +/** + * Creates a signal: waiting on it resolves the next time it is notified. + * + * @returns {{ wait: () => Promise, notify: () => void }} + */ +export const createSignal = () => { + let next = Promise.withResolvers(); + + return { + /** + * Waits for the next notification. + * + * @returns {Promise} + */ + wait: () => next.promise, + + /** + * Wakes what waits. + */ + notify: () => { + next.resolve(); + next = Promise.withResolvers(); + }, + }; +};