perf: stream a generator's output to the generators that read it

.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.

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.

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<Array<{ value: number }>>} reader
+ * @param {number} [count]
+ * @returns {Promise<number[]>}
+ */
+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);
+    });
   });
 });

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<number>, 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<void>}
+ */
+export async function collectGarbage(done) {
+  for (let i = 0; i < 10 && !done(); i++) {
+    await new Promise(resolve => setImmediate(resolve));
+    gc();
+  }
+}

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 }]);
+  });
 });