diff --git a/.changeset/stream-input.md b/.changeset/stream-input.md new file mode 100644 index 000000000..bd84fc1de --- /dev/null +++ b/.changeset/stream-input.md @@ -0,0 +1,5 @@ +--- +'@doc-kit/core': minor +--- + +Generators can set `streamsInput: true` to receive their dependency's output as an async iterable of chunks as they are produced, instead of waiting for all of it. The cache runs a streamed generator once for every reader, and lets a chunk go once all of them have read it. diff --git a/docs/creating-generators.md b/docs/creating-generators.md index c2912d335..8e15bb89f 100644 --- a/docs/creating-generators.md +++ b/docs/creating-generators.md @@ -330,9 +330,36 @@ 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 with `streamsInput` 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`, `input` is instead an +async iterable of the dependency's chunks, read as they are produced, which +`worker.stream` accepts directly: + +```mjs displayName="index.mjs" +import { processChunk } from './generate.mjs'; + +export default { + name: 'per-chunk-generator', + dependsOn: '@doc-kit/core/metadata', + hasParallelProcessor: true, + streamsInput: true, + async *generate(input, worker) { + yield* worker.stream(input, {}); + }, + processChunk, +}; +``` + +- A dependency that does not stream arrives as a single chunk. +- With one thread, `input` is still the collected array, so accept both. +- `input` can be iterated only once. + ### 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..b6106f345 100644 --- a/packages/core/src/__tests__/caching.test.mjs +++ b/packages/core/src/__tests__/caching.test.mjs @@ -1,7 +1,35 @@ -import { deepStrictEqual, strictEqual } from 'node:assert'; +import { deepStrictEqual, strictEqual, throws } from 'node:assert'; import { describe, it } from 'node:test'; import { collectAsyncGenerator, createCache } from '../caching.mjs'; +import { + collectGarbage, + trackReleases, +} from './fixtures/garbage-collection.mjs'; + +/** + * 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; +} describe('caching', () => { describe('collectAsyncGenerator', () => { @@ -106,5 +134,123 @@ describe('caching', () => { await cache.consume('a'); strictEqual(cache.has('a'), false); }); + + 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(stream => Array.fromAsync(stream)), + cache.consume('a'), + ]); + + deepStrictEqual(streamed, [[1, 2], [3]]); + deepStrictEqual(collected, [1, 2, 3]); + strictEqual(iterations, 1); + }); + + it('should stream a non-streaming result as a single chunk', async () => { + const cache = createCache(); + + cache.store('a', Promise.resolve([1, 2])); + + const stream = await cache.stream('a'); + + deepStrictEqual(await Array.fromAsync(stream), [[1, 2]]); + }); + + 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'); + + deepStrictEqual(await Array.fromAsync(stream), [[1]]); + throws(() => stream[Symbol.asyncIterator](), /can only be read once/); + }); + + it('should run a stream before it is read, and count it as a read', async () => { + const cache = createCache(); + + let produced = 0; + + cache.populateConsumerCounts(['a'], () => undefined); + + async function* gen() { + for (const value of [1, 2, 3]) { + produced++; + + yield [value]; + } + } + + 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); + + await new Promise(resolve => setImmediate(resolve)); + + // Every chunk is produced before the reader reads any + strictEqual(produced, 3); + deepStrictEqual(await Array.fromAsync(stream), [[1], [2], [3]]); + }); + + it('should let go of a chunk once every streaming reader has read it, while it still runs', async () => { + const cache = createCache(); + const gate = Promise.withResolvers(); + const { released, chunkOf } = trackReleases(); + + // `x` and `y` both read `a` + cache.populateConsumerCounts(['x', 'y'], name => + name === 'a' ? undefined : 'a' + ); + + async function* gen() { + yield chunkOf(1); + yield chunkOf(2); + await gate.promise; + } + + cache.store('a', gen()); + + // The second is the last read: no more readers can start + const first = (await cache.stream('a'))[Symbol.asyncIterator](); + const second = (await cache.stream('a'))[Symbol.asyncIterator](); + + deepStrictEqual(await readValues(first, 2), [1, 2]); + + // `second` has yet to read the first chunk + await collectGarbage(() => released.has(1)); + + strictEqual(released.has(1), false); + + deepStrictEqual(await readValues(second, 2), [1, 2]); + + await collectGarbage(() => released.has(1)); + + strictEqual(released.has(1), true); + + gate.resolve(); + + await readValues(first); + await readValues(second); + }); }); }); diff --git a/packages/core/src/__tests__/fixtures/garbage-collection.mjs b/packages/core/src/__tests__/fixtures/garbage-collection.mjs new file mode 100644 index 000000000..d83f5bc06 --- /dev/null +++ b/packages/core/src/__tests__/fixtures/garbage-collection.mjs @@ -0,0 +1,47 @@ +import { setFlagsFromString } from 'node:v8'; +import { runInNewContext } from 'node:vm'; + +setFlagsFromString('--expose-gc'); + +const gc = runInNewContext('gc'); + +/** + * Makes chunks of one item each that tell when they are let go. + * + * @returns {{ released: Set, chunkOf: (value: number) => Array<{ value: number }> }} + */ +export 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} + */ +export async function collectGarbage(done) { + for (let i = 0; i < 10 && !done(); i++) { + await new Promise(resolve => setImmediate(resolve)); + gc(); + } +} diff --git a/packages/core/src/__tests__/generators.test.mjs b/packages/core/src/__tests__/generators.test.mjs index d3be80da3..457393439 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; })(); +// Opened by `stream-reader` once it has read `gated-source`'s first chunk +const gate = Promise.withResolvers(); + // Synthetic generators keyed by specifier (any string works as a specifier) const syntheticGenerators = { // Root streaming generator (no dependency) @@ -99,6 +102,52 @@ const syntheticGenerators = { return { all: input }; }, }, + // A source that only yields its second chunk once its streaming reader has + // read the first + 'gated-source': { + name: 'gated-source', + hasParallelProcessor: true, + generate: async () => + (async function* () { + yield [1]; + await gate.promise; + yield [2]; + })(), + }, + 'stream-reader': { + name: 'stream-reader', + dependsOn: 'gated-source', + streamsInput: true, + async *generate(input) { + for await (const chunk of input) { + gate.resolve(); + + yield [{ streamed: chunk }]; + } + }, + }, + // A source that fails before yielding anything, read as a stream + 'failing-source': { + name: 'failing-source', + hasParallelProcessor: true, + generate: async () => { + throw new Error('source failed'); + }, + }, + 'failing-reader': { + name: 'failing-reader', + dependsOn: 'failing-source', + streamsInput: true, + async *generate(input) { + yield* input; + }, + }, + 'shape-reader': { + name: 'shape-reader', + dependsOn: 'metadata', + streamsInput: true, + generate: async input => ({ array: Array.isArray(input) }), + }, }; mock.module('../generators/loader.mjs', { @@ -245,4 +294,57 @@ describe('createGenerator orchestration', () => { { name: 'gen-a', markdown, ran: ['ast', 'metadata'] }, ]); }); + + it( + 'streams a dependency to a generator that streams its input', + // The source waits for its first chunk to be read, so a reader that only + // got its input once the source completed would never get it + { timeout: 10_000 }, + async () => { + const { runGenerators } = createGenerator(); + + const results = await runGenerators({ + target: ['stream-reader'], + threads: 2, + }); + + assert.deepStrictEqual(results, [[{ streamed: [1] }, { streamed: [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: ['failing-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..ce533e104 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 { createTee } from './utils/tee.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, buffering 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 tees of the results that stream, and those of them collected into + // one array, shared by the consumers that collect them + /** @type {Map>} */ + const tees = new Map(); /** @type {Map>} */ const collections = new Map(); @@ -66,27 +84,39 @@ 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 tee 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 teeOf = (name, generator) => { + if (!tees.has(name)) { + cachingLogger.debug(`Starting the tee of "${name}"`); - if (!hasKey) { - collections.set(key, collectAsyncGenerator(generator)); + tees.set(name, createTee(generator)); } - cachingLogger.debug( - hasKey - ? `Using cached collection for "${key}"` - : `Starting collection for "${key}"` - ); + return tees.get(name); + }; + + /** + * Counts a read, evicting the result 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 + tees.get(name)?.forget(); + tees.delete(name); + collections.delete(name); - return collections.get(key); + cachingLogger.debug(`Evicted "${name}" after final consumer`); + } }; return { @@ -165,19 +195,67 @@ 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(teeOf(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 tee 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 = teeOf(name, result).read(); + } else { + 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..06b497d54 100644 --- a/packages/core/src/generators.mjs +++ b/packages/core/src/generators.mjs @@ -28,17 +28,18 @@ const createGenerator = () => { let pool; /** - * Gets the collected input from a dependency generator. + * Gets the input from a dependency generator, collected or streamed. * * @param {string | undefined} dependsOn - Resolved dependency specifier + * @param {boolean} [streamed] - Whether to stream its chunks * @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 +58,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 +76,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 +97,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 + // A streaming generator returns before it has produced anything, so + // it isn't complete here 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..7bac2a412 100644 --- a/packages/core/src/generators/types.d.ts +++ b/packages/core/src/generators/types.d.ts @@ -13,16 +13,16 @@ declare global { */ export interface ParallelWorker { /** - * Processes items in parallel across worker threads and yields results - * as each chunk completes. Enables downstream processing to begin - * while upstream chunks are still being processed. + * Processes items in parallel across worker threads and yields each + * chunk's results in submission order. 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, or a streamed input's batches of 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 +70,13 @@ declare global { hasParallelProcessor?: boolean; + /** + * Receive the dependency's output as an async iterable of its chunks, as + * they are produced, instead of one collected array. With one thread, the + * input is still a collected array, 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__/tee.test.mjs b/packages/core/src/utils/__tests__/tee.test.mjs new file mode 100644 index 000000000..08a5ba11f --- /dev/null +++ b/packages/core/src/utils/__tests__/tee.test.mjs @@ -0,0 +1,142 @@ +import { deepStrictEqual, rejects } from 'node:assert'; +import { describe, it } from 'node:test'; + +import { + collectGarbage, + trackReleases, +} from '../../__tests__/fixtures/garbage-collection.mjs'; +import { createTee } from '../tee.mjs'; + +describe('createTee', () => { + it('should let go of the chunks every reader has read while it still runs', async () => { + const { released, chunkOf } = trackReleases(); + const gate = Promise.withResolvers(); + + async function* gen() { + for (let value = 0; value < 4; value++) { + yield chunkOf(value); + } + + await gate.promise; + + yield chunkOf(4); + } + + const tee = createTee(gen()); + const readers = [tee.read(), tee.read()]; + + tee.forget(); + + for (const reader of readers) { + for (let i = 0; i < 4; i++) { + await reader.next(); + } + } + + // The last chunk read may still be held where each reader is suspended + await collectGarbage(() => [0, 1, 2].every(value => released.has(value))); + + deepStrictEqual( + [0, 1, 2].filter(value => released.has(value)), + [0, 1, 2] + ); + + gate.resolve(); + + for (const reader of readers) { + await Array.fromAsync(reader); + } + + // Once its readers are done, so is the last chunk + await collectGarbage(() => released.has(4)); + + deepStrictEqual(released.size, 5); + }); + + 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 tee = createTee(gen()); + + const early = Array.fromAsync(tee.read()); + + await new Promise(resolve => setImmediate(resolve)); + gate.resolve(); + + const late = Array.fromAsync(tee.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 tee = createTee(gen()); + + await tee.ended; + + // With no reader left, every chunk goes + tee.forget(); + + await rejects(Array.fromAsync(tee.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('boom'); + } + + const tee = createTee(gen()); + const chunks = []; + + await rejects(tee.ended, /boom/); + + await rejects(async () => { + for await (const chunk of tee.read()) { + chunks.push(chunk); + } + }, /boom/); + + deepStrictEqual(chunks, [[1]]); + }); + + it('should not report a failure as unhandled when nothing reads it', async () => { + const unhandled = []; + const onUnhandled = error => unhandled.push(error); + + async function* gen() { + yield [1]; + + throw new Error('boom'); + } + + process.on('unhandledRejection', onUnhandled); + + try { + createTee(gen()); + + // A rejection left unhandled is reported once the microtasks have run + await new Promise(resolve => setImmediate(resolve)); + } finally { + process.off('unhandledRejection', onUnhandled); + } + + deepStrictEqual(unhandled, []); + }); +}); diff --git a/packages/core/src/utils/tee.mjs b/packages/core/src/utils/tee.mjs new file mode 100644 index 000000000..18889569a --- /dev/null +++ b/packages/core/src/utils/tee.mjs @@ -0,0 +1,101 @@ +'use strict'; + +/** + * Reads the chunks linked from `node`, then throws if the generator failed. + * A chunk is held only by the link before it, so one every reader has passed + * is let go. + * + * @template T + * @param {Promise<{ chunk: T[], next: Promise } | undefined> | undefined} node + * @returns {AsyncGenerator} + */ +async function* readFrom(node) { + if (!node) { + throw new Error('Chunk 0 was let go before this reader read it'); + } + + for (;;) { + const link = await node; + + if (!link) { + return; + } + + node = link.next; + + yield link.chunk; + } +} + +/** + * Splits an async generator into any number of readers, like + * `ReadableStream.prototype.tee()`: an async generator can only be iterated + * once, while several generators may read the same dependency. + * + * It runs the generator to its end, buffering its chunks so that each reader + * can 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 createTee = generator => { + let tail = Promise.withResolvers(); + let head = tail.promise; + + /** + * Runs the generator to its end, linking each chunk after the last, whether + * or not its readers keep up. + */ + const run = async () => { + try { + for await (const chunk of generator) { + const next = Promise.withResolvers(); + + tail.resolve({ chunk, next: next.promise }); + tail = next; + } + + tail.resolve(); + } catch (error) { + // Readers get it as they read; a tee with none mustn't report it + tail.promise.catch(() => {}); + tail.reject(error); + + throw error; + } finally { + // Let the finished generator, and all it holds, go + generator = undefined; + } + }; + + const ended = run(); + + 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: () => readFrom(head), + + /** + * Says no more readers will start, so chunks every reader has read can go. + */ + forget: () => { + head = undefined; + }, + }; +};