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/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(); + }, + }; +}; 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(); + }, + }; +};