stream: speed up async iteration of Readable

Replace the async generator backing Symbol.asyncIterator with a
hand-rolled iterator. The generator machinery costs several extra
promise allocations and microtask hops per chunk: yield awaits the
yielded value and resolves the pending request through separate
promises. Buffered chunks are now delivered as an already-resolved
promise, one microtask sooner than before.

Thenable chunks are still awaited before delivery, requests received
while a next() is outstanding are queued, and return()/throw() before
the first next() complete the iterator without touching the stream.

The earlier delivery is observable by code racing an abort against
the first chunk. The flatMap AbortSignal test relied on such a race;
it is reworked to abort deterministically while two mappers are in
flight, asserting the concurrency limit, in-flight cancellation and
rejection, without depending on delivery timing or timers.

streams/readable-async-iterator.js sync='yes': +32.59% (***)
streams/readable-async-iterator.js sync='no': +9.84% (***)

Assisted-by: Claude Fable 5
Signed-off-by: Matteo Collina <matteo.collina@gmail.com>

Author: mcollina

lib/internal/streams/readable.js

@@ -30,6 +30,9 @@ const {
   ObjectKeys,
   ObjectSetPrototypeOf,
   Promise,
+  PromisePrototypeThen,
+  PromiseReject,
+  PromiseResolve,
   ReflectApply,
   SafeSet,
   Symbol,
@@ -1382,10 +1385,21 @@ function streamToAsyncIterator(stream, options) {
   return iter;
 }
 
-async function* createAsyncIterator(stream, options) {
+// Async iterator over a Readable. Requests received while another is
+// outstanding are queued and processed in order.
+function createAsyncIterator(stream, options) {
   let callback = nop;
-
-  function next(resolve) {
+  let error;            // undefined: active, null: ended cleanly, else: Error
+  let started = false;
+  let completed = false;
+  let inFlight = false; // An asynchronous request is outstanding
+  let queue = null;     // Requests received while inFlight
+  let cleanup;
+
+  // Used both as the 'readable' listener (where `this === stream`) and
+  // as a promise executor storing the resolver that wakes up a pending
+  // pump().
+  function wakeup(resolve) {
     if (this === stream) {
       callback();
       callback = nop;
@@ -1394,32 +1408,23 @@ async function* createAsyncIterator(stream, options) {
     }
   }
 
-  stream.on('readable', next);
+  function start() {
+    started = true;
 
-  let error;
-  const cleanup = eos(stream, { writable: false }, (err) => {
-    error = err ? aggregateTwoErrors(error, err) : null;
-    callback();
-    callback = nop;
-  });
+    stream.on('readable', wakeup);
+
+    cleanup = eos(stream, { writable: false }, (err) => {
+      error = err ? aggregateTwoErrors(error, err) : null;
+      callback();
+      callback = nop;
+    });
+  }
+
+  // Complete the iterator and either destroy the stream or detach
+  // from it.
+  function finalize() {
+    completed = true;
 
-  try {
-    while (true) {
-      const chunk = stream.destroyed ? null : stream.read();
-      if (chunk !== null) {
-        yield chunk;
-      } else if (error) {
-        throw error;
-      } else if (error === null) {
-        return;
-      } else {
-        await new Promise(next);
-      }
-    }
-  } catch (err) {
-    error = aggregateTwoErrors(error, err);
-    throw error;
-  } finally {
     const preserveHalfOpenDuplex =
       error === null &&
       stream.allowHalfOpen === true &&
@@ -1433,10 +1438,174 @@ async function* createAsyncIterator(stream, options) {
     ) {
       destroyImpl.destroyer(stream, null);
     } else {
-      stream.off('readable', next);
+      stream.off('readable', wakeup);
       cleanup();
     }
   }
+
+  function settleError(err, reject) {
+    error = aggregateTwoErrors(error, err);
+    finalize();
+    reject(error);
+  }
+
+  function drain() {
+    while (!inFlight && queue.length > 0) {
+      const req = queue.shift();
+      if (req.type === 'next') {
+        processNext(req.resolve, req.reject);
+      } else if (req.type === 'return') {
+        processReturn(req.value, req.resolve);
+      } else {
+        processThrow(req.value, req.reject);
+      }
+    }
+  }
+
+  // Thenable chunks are unwrapped before delivery; a rejection tears
+  // down the iterator and the stream.
+  function onChunkFulfilled(value) {
+    inFlight = false;
+    if (queue !== null) drain();
+    return { done: false, value };
+  }
+
+  function onChunkRejected(err) {
+    inFlight = false;
+    error = aggregateTwoErrors(error, err);
+    finalize();
+    if (queue !== null) drain();
+    throw error;
+  }
+
+  // Runs with inFlight === true; settles the request and hands over to
+  // any requests that queued up behind it.
+  function pump(resolve, reject) {
+    const chunk = stream.destroyed ? null : stream.read();
+    if (chunk !== null) {
+      if (typeof chunk.then === 'function') {
+        PromisePrototypeThen(PromiseResolve(chunk), (value) => {
+          inFlight = false;
+          resolve({ done: false, value });
+          if (queue !== null) drain();
+        }, (err) => {
+          inFlight = false;
+          settleError(err, reject);
+          if (queue !== null) drain();
+        });
+        return;
+      }
+      inFlight = false;
+      resolve({ done: false, value: chunk });
+      if (queue !== null) drain();
+    } else if (error) {
+      inFlight = false;
+      settleError(error, reject);
+      if (queue !== null) drain();
+    } else if (error === null) {
+      inFlight = false;
+      finalize();
+      resolve({ done: true, value: undefined });
+      if (queue !== null) drain();
+    } else {
+      // No data buffered yet; wait for 'readable' or end-of-stream and
+      // retry.
+      PromisePrototypeThen(new Promise(wakeup), () => pump(resolve, reject));
+    }
+  }
+
+  function processNext(resolve, reject) {
+    if (completed) {
+      resolve({ done: true, value: undefined });
+      return;
+    }
+    if (!started) start();
+    inFlight = true;
+    pump(resolve, reject);
+  }
+
+  function processReturn(value, resolve) {
+    if (!completed) {
+      if (started) {
+        finalize();
+      } else {
+        // Never started: complete without touching the stream.
+        completed = true;
+      }
+    }
+    resolve({ done: true, value });
+  }
+
+  function processThrow(err, reject) {
+    if (completed || !started) {
+      completed = true;
+      reject(err);
+      return;
+    }
+    settleError(err, reject);
+  }
+
+  return {
+    next() {
+      if (!inFlight && !completed) {
+        if (!started) start();
+        // Fast path: a chunk is already buffered.
+        const chunk = stream.destroyed ? null : stream.read();
+        if (chunk !== null) {
+          if (typeof chunk.then === 'function') {
+            inFlight = true;
+            return PromisePrototypeThen(
+              PromiseResolve(chunk), onChunkFulfilled, onChunkRejected);
+          }
+          return PromiseResolve({ done: false, value: chunk });
+        }
+        if (error) {
+          finalize();
+          return PromiseReject(error);
+        }
+        if (error === null) {
+          finalize();
+          return PromiseResolve({ done: true, value: undefined });
+        }
+        // No data buffered yet; wait for 'readable' or end-of-stream.
+        inFlight = true;
+        return new Promise((resolve, reject) => {
+          PromisePrototypeThen(new Promise(wakeup), () => pump(resolve, reject));
+        });
+      }
+      return new Promise((resolve, reject) => {
+        if (inFlight) {
+          queue ??= [];
+          queue.push({ type: 'next', value: undefined, resolve, reject });
+        } else {
+          resolve({ done: true, value: undefined });
+        }
+      });
+    },
+    return(value) {
+      return new Promise((resolve, reject) => {
+        if (inFlight) {
+          queue ??= [];
+          queue.push({ type: 'return', value, resolve, reject });
+        } else {
+          processReturn(value, resolve);
+        }
+      });
+    },
+    throw(err) {
+      return new Promise((resolve, reject) => {
+        if (inFlight) {
+          queue ??= [];
+          queue.push({ type: 'throw', value: err, resolve, reject });
+        } else {
+          processThrow(err, reject);
+        }
+      });
+    },
+    [SymbolAsyncIterator]() {
+      return this;
+    },
+  };
 }
 
 let composeImpl;

test/parallel/test-stream-flatMap.js

@@ -72,10 +72,23 @@ function oneTo5() {
 
 {
   // Concurrency + AbortSignal
+  // Two mappers are started concurrently and block until their signal
+  // is aborted. Aborting while both are in flight must cancel them and
+  // reject the iteration, without ever starting a third mapper.
   const ac = new AbortController();
-  const stream = oneTo5().flatMap(common.mustNotCall(async (_, { signal }) => {
-    await setTimeout(100, { signal });
-  }), { signal: ac.signal, concurrency: 2 });
+  const stream = oneTo5().flatMap(common.mustCall(async (x, { signal }) => {
+    if (x === 2) {
+      // Both mappers allowed by `concurrency` are now in flight.
+      ac.abort();
+    }
+    await new Promise((resolve, reject) => {
+      if (signal.aborted) {
+        reject(signal.reason);
+        return;
+      }
+      signal.addEventListener('abort', () => reject(signal.reason), { once: true });
+    });
+  }, 2), { signal: ac.signal, concurrency: 2 });
   // pump
   assert.rejects(async () => {
     for await (const item of stream) {
@@ -85,10 +98,6 @@ function oneTo5() {
   }, {
     name: 'AbortError',
   }).then(common.mustCall());
-
-  queueMicrotask(() => {
-    ac.abort();
-  });
 }
 
 {

test/parallel/test-stream-readable-async-iterators.js

@@ -869,5 +869,99 @@ async function tests() {
   }));
 }
 
+{
+  // Thenable chunks are awaited before delivery.
+  (async () => {
+    const r = new Readable({ objectMode: true, read() {} });
+    r.push(Promise.resolve('unwrapped'));
+    r.push(null);
+
+    const it = r[Symbol.asyncIterator]();
+    const { value, done } = await it.next();
+    assert.strictEqual(done, false);
+    assert.strictEqual(value, 'unwrapped');
+  })().then(common.mustCall());
+}
+
+{
+  // A rejected thenable chunk tears down the iterator and the stream.
+  (async () => {
+    const r = new Readable({ objectMode: true, read() {} });
+    const rejected = Promise.reject(new Error('kaboom'));
+    rejected.catch(() => {});
+    r.push(rejected);
+    r.push(null);
+
+    const it = r[Symbol.asyncIterator]();
+    await assert.rejects(it.next(), { message: 'kaboom' });
+    assert.strictEqual((await it.next()).done, true);
+    assert.strictEqual(r.destroyed, true);
+  })().then(common.mustCall());
+}
+
+{
+  // throw() rejects with the passed error, destroys the stream and
+  // completes the iterator.
+  (async () => {
+    const r = new Readable({ objectMode: true, read() {} });
+    r.push('a');
+
+    const it = r[Symbol.asyncIterator]();
+    assert.strictEqual((await it.next()).value, 'a');
+    await assert.rejects(it.throw(new Error('kaboom')), { message: 'kaboom' });
+    assert.strictEqual(r.destroyed, true);
+    assert.strictEqual((await it.next()).done, true);
+  })().then(common.mustCall());
+}
+
+{
+  // throw() before the first next() completes the iterator without
+  // touching the stream.
+  (async () => {
+    const r = new Readable({ objectMode: true, read() {} });
+    const it = r[Symbol.asyncIterator]();
+    await assert.rejects(it.throw(new Error('kaboom')), { message: 'kaboom' });
+    assert.strictEqual(r.destroyed, false);
+    assert.strictEqual(r.listenerCount('readable'), 0);
+    assert.strictEqual((await it.next()).done, true);
+    r.destroy();
+  })().then(common.mustCall());
+}
+
+{
+  // Concurrent next() calls while waiting for data are served in order.
+  (async () => {
+    const r = new Readable({ objectMode: true, read() {} });
+    const it = r[Symbol.asyncIterator]();
+
+    const p1 = it.next();
+    const p2 = it.next();
+    r.push('a');
+    r.push('b');
+    r.push(null);
+
+    assert.strictEqual((await p1).value, 'a');
+    assert.strictEqual((await p2).value, 'b');
+    assert.strictEqual((await it.next()).done, true);
+  })().then(common.mustCall());
+}
+
+{
+  // return() while a next() is pending is processed after the pending
+  // next() settles, and destroys the stream.
+  (async () => {
+    const r = new Readable({ objectMode: true, read() {} });
+    const it = r[Symbol.asyncIterator]();
+
+    const p1 = it.next();
+    const p2 = it.return();
+    r.push('a');
+
+    assert.strictEqual((await p1).value, 'a');
+    assert.deepStrictEqual(await p2, { done: true, value: undefined });
+    assert.strictEqual(r.destroyed, true);
+  })().then(common.mustCall());
+}
+
 // To avoid missing some tests if a promise does not resolve
 tests().then(common.mustCall());