stream: speed up async iteration of Readable

lib/internal/streams/readable.js

@@ -23,13 +23,18 @@
 
 const {
   ArrayPrototypeIndexOf,
+  AsyncIteratorPrototype,
+  FunctionPrototypeCall,
   NumberIsInteger,
   NumberIsNaN,
   NumberParseInt,
   ObjectDefineProperties,
   ObjectKeys,
   ObjectSetPrototypeOf,
   Promise,
+  PromisePrototypeThen,
+  PromiseReject,
+  PromiseResolve,
   ReflectApply,
   SafeSet,
   Symbol,
@@ -100,6 +105,7 @@ const FastBuffer = Buffer[SymbolSpecies];
 
 const { StringDecoder } = require('string_decoder');
 const from = require('internal/streams/from');
+const FixedQueue = require('internal/fixed_queue');
 
 ObjectSetPrototypeOf(Readable.prototype, Stream.prototype);
 ObjectSetPrototypeOf(Readable, Stream);
@@ -1382,10 +1388,22 @@ 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 draining = false;
+  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 +1412,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 +1442,191 @@ 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() {
+    // Requests settled synchronously call back into drain(); the guard
+    // keeps a single loop going instead of recursing once per request.
+    if (draining) {
+      return;
+    }
+    draining = true;
+    try {
+      while (!inFlight && !queue.isEmpty()) {
+        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);
+        }
+      }
+    } finally {
+      draining = false;
+    }
+  }
+
+  // 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) {
+      // Read `then` only once so that a getter cannot observe (or throw
+      // on) a second access.
+      const then = chunk.then;
+      if (typeof then === 'function') {
+        FunctionPrototypeCall(then, 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 {
+    __proto__: AsyncIteratorPrototype,
+    next() {
+      if (!inFlight && !completed) {
+        if (!started) start();
+        // Fast path: a chunk is already buffered.
+        const chunk = stream.destroyed ? null : stream.read();
+        if (chunk !== null) {
+          // Read `then` only once so that a getter cannot observe (or
+          // throw on) a second access.
+          const then = chunk.then;
+          if (typeof then === 'function') {
+            inFlight = true;
+            return FunctionPrototypeCall(
+              then, 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 ??= new FixedQueue();
+          queue.push({ __proto__: null, type: 'next', value: undefined, resolve, reject });
+        } else {
+          resolve({ done: true, value: undefined });
+        }
+      });
+    },
+    return(value) {
+      return new Promise((resolve, reject) => {
+        if (inFlight) {
+          queue ??= new FixedQueue();
+          queue.push({ __proto__: null, type: 'return', value, resolve, reject });
+        } else {
+          processReturn(value, resolve);
+        }
+      });
+    },
+    throw(err) {
+      return new Promise((resolve, reject) => {
+        if (inFlight) {
+          queue ??= new FixedQueue();
+          queue.push({ __proto__: null, 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();
+    }
+    const { promise, reject } = Promise.withResolvers();
+    if (signal.aborted) {
+      reject(signal.reason);
+    }
+    signal.addEventListener('abort', () => reject(signal.reason), { once: true });
+    // Promise is expected to reject.
+    await promise;
+  }, 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();
-  });
 }
 
 {