stream: read async sources more cheaply in stream/iter pipeTo()

from() reads an async iterable source through a layer that lets a
cancellation of the normalization reject a pending read: for every
batch, it waits with a PromiseWithResolvers() and three closures, and
the normalizer handles the source's result in a second reaction.

pipeTo() never cancels the normalization while a read is pending: it
only calls return() after an error writing a batch. Without transforms
or a signal, have it read through a method of the iterator returned by
from() that reads the source without waiting for a cancellation, and
handles the source's result in a single reaction. The source's results
are checked as before, a write error still closes the source, and an
error reading the source still does not.

Piping an async generator yielding 16-byte chunks, one per batch, is
about 1.5 times faster, with one promise and 620 bytes less per batch.

Assisted-by: OpenCode

Author: jasnell

lib/internal/streams/iter/from.js

@@ -144,9 +144,44 @@ function waitForNormalization(value, context) {
 // in createSyncSourceNormalizer().
 const kNextSyncBatch = Symbol('kNextSyncBatch');
 
+// The method of an iterator returned by from() for an async iterable that
+// reads like next() when nothing can cancel the normalization while a read
+// is pending, for pipeTo(): see createAsyncSourceNormalizer().
+const kNextUncancellable = Symbol('kNextUncancellable');
+
+// Methods of the iterator returned by yieldNormalizationAbortable(), used by
+// createAsyncSourceNormalizer() for kNextUncancellable.
+const kReadUncancellable = Symbol('kReadUncancellable');
+const kToIterResult = Symbol('kToIterResult');
+const kReadFailed = Symbol('kReadFailed');
+// The result kReadUncancellable gives once the source is done.
+const kDoneSourceResult = ObjectFreeze({ __proto__: null, done: true, value: undefined });
+
 function createNormalizationIterator(createIterator) {
   const context = createNormalizationContext();
   const iterator = createIterator(context);
+  if (iterator[kNextUncancellable] !== undefined) {
+    return ObjectSetPrototypeOf({
+      next(value) {
+        return FunctionPrototypeCall(iterator.next, iterator, value);
+      },
+      return(value) {
+        cancelNormalization(
+          context, lazyDOMException('Aborted', 'AbortError'));
+        return FunctionPrototypeCall(iterator.return, iterator, value);
+      },
+      throw(error) {
+        cancelNormalization(context, error, true);
+        return FunctionPrototypeCall(iterator.throw, iterator, error);
+      },
+      [kNextUncancellable]() {
+        return iterator[kNextUncancellable]();
+      },
+      [SymbolAsyncIterator]() {
+        return this;
+      },
+    }, null);
+  }
   if (iterator[kNextSyncBatch] !== undefined) {
     return ObjectSetPrototypeOf({
       next(value) {
@@ -468,6 +503,15 @@ function yieldNormalizationAbortable(source, context) {
         return new IterResult(false, value);
       }
 
+      function onUncancellableResult(result) {
+        try {
+          return toIterResult(result);
+        } catch (error) {
+          reading = false;
+          throw error;
+        }
+      }
+
       // Reject a pending next(), closing the source first if the
       // normalization has been cancelled.
       function rejectNext(reject, error) {
@@ -535,6 +579,30 @@ function yieldNormalizationAbortable(source, context) {
             });
           return promise;
         },
+        // Like next(), when nothing cancels the normalization while the read
+        // is pending, in two parts so that the caller handles the result in
+        // the same reaction: kReadUncancellable starts the read, returning a
+        // promise for the source's result, and the caller passes that result
+        // to kToIterResult, or calls kReadFailed if the promise rejects.
+        [kReadUncancellable]() {
+          if (completed) {
+            return PromiseResolve(kDoneSourceResult);
+          }
+          if (context.cancelled) return PromiseReject(context.reason);
+          reading = true;
+          let next;
+          try {
+            next = FunctionPrototypeCall(nextMethod, iterator);
+          } catch (error) {
+            reading = false;
+            return PromiseReject(error);
+          }
+          return PromiseResolve(next);
+        },
+        [kToIterResult]: onUncancellableResult,
+        [kReadFailed]() {
+          reading = false;
+        },
         async return(value) {
           await closeSource(
             context.suppressCleanup || (context.cancelled && reading));
@@ -723,6 +791,10 @@ function createAsyncSourceNormalizer(source, context) {
   // oversized batch, or a normalizeAsyncSourceValue() generator.
   let boundedBatches = null;
   let valueBatches = null;
+  // Whether the source is read with kNextUncancellable: set by the
+  // kNextUncancellable method, for callers that never cancel the
+  // normalization while a read is pending.
+  let uncancellable = false;
   const { run, settled } = createOperationQueue();
 
   function finish(result) {
@@ -784,7 +856,28 @@ function createAsyncSourceNormalizer(source, context) {
     return finish(new IterResult(false, result.value));
   }
 
+  // A kReadUncancellable result: handled as the result of the source's
+  // next() is, in the same reaction.
+  function onUncancellableResult(result) {
+    let iterResult;
+    try {
+      iterResult = iterator[kToIterResult](result);
+    } catch (error) {
+      return fail(error);
+    }
+    return onSourceResult(iterResult);
+  }
+
+  function onUncancellableError(error) {
+    iterator[kReadFailed]();
+    return fail(error);
+  }
+
   function pullSource() {
+    if (uncancellable) {
+      return PromisePrototypeThen(iterator[kReadUncancellable](),
+                                  onUncancellableResult, onUncancellableError);
+    }
     return PromisePrototypeThen(iterator.next(), onSourceResult, fail);
   }
 
@@ -866,6 +959,10 @@ function createAsyncSourceNormalizer(source, context) {
     next() { return run(doNext); },
     return(value) { return run(doReturn, value); },
     throw(error) { return run(doThrow, error); },
+    [kNextUncancellable]() {
+      uncancellable = true;
+      return run(doNext);
+    },
   }, null);
 }
 
@@ -1344,6 +1441,7 @@ module.exports = {
   isSyncIterable,
   isUint8ArrayBatch,
   kNextSyncBatch,
+  kNextUncancellable,
   normalizeAsyncSource,
   normalizeAsyncValue,
   normalizeSyncSource,

lib/internal/streams/iter/pull.js

@@ -53,6 +53,7 @@ const {
   isAsyncIterable,
   isUint8ArrayBatch,
   kNextSyncBatch,
+  kNextUncancellable,
 } = require('internal/streams/iter/from');
 
 const {
@@ -1737,6 +1738,27 @@ async function pipeTo(source, ...args) {
     }
   }
 
+  // Write the batches of an async source as the for await...of loop below
+  // does. Nothing cancels the normalization while a batch is read: return()
+  // is only called after an error writing one.
+  async function pipeAsyncSource(iterator) {
+    for (;;) {
+      const result = await iterator[kNextUncancellable]();
+      if (result.done) return;
+      try {
+        const p = writeBatch(result.value);
+        if (p) await p;
+      } catch (error) {
+        try {
+          await iterator.return();
+        } catch {
+          // The error writing the batch is thrown.
+        }
+        throw error;
+      }
+    }
+  }
+
   function writeBatch(batch) {
     // Single chunk, the common case: check the view around writeSync()
     // without allocating a batch entry, and create one only to fall back to
@@ -1798,6 +1820,8 @@ async function pipeTo(source, ...args) {
         }
       } else if (normalized[kNextSyncBatch] !== undefined) {
         await pipeSyncSource(normalized);
+      } else if (normalized[kNextUncancellable] !== undefined) {
+        await pipeAsyncSource(normalized);
       } else {
         for await (const batch of normalized) {
           const p = writeBatch(batch);

test/parallel/test-stream-iter-pipeto.js

@@ -391,6 +391,63 @@ async function testPipeToSyncIterableWriteError() {
   }), error);
 }
 
+// pipeTo() reads async iterables without waiting for a cancellation, since
+// nothing cancels them while a read is pending. An error writing a batch
+// must still close the source, ignoring an error closing it, an error
+// reading the source must not close it, and the source's results are still
+// checked.
+async function testPipeToAsyncIterableErrors() {
+  for (const returnRejects of [false, true]) {
+    const error = new Error('write');
+    let closed = false;
+    const source = {
+      [Symbol.asyncIterator]() {
+        let i = 0;
+        return {
+          async next() {
+            return { done: false, value: [new Uint8Array([i++])] };
+          },
+          async return() {
+            closed = true;
+            if (returnRejects) throw new Error('return');
+            return { done: true };
+          },
+        };
+      },
+    };
+    await assert.rejects(pipeTo(source, {
+      write: common.mustNotCall(),
+      writeSync: common.mustCall(() => { throw error; }),
+      fail: common.mustCall((reason) => assert.strictEqual(reason, error)),
+    }), error);
+    assert.strictEqual(closed, true);
+  }
+
+  const error = new Error('source');
+  await assert.rejects(pipeTo({
+    [Symbol.asyncIterator]() {
+      return {
+        next() { return Promise.reject(error); },
+        return: common.mustNotCall(),
+      };
+    },
+  }, {
+    write: common.mustNotCall(),
+    writeSync: common.mustNotCall(),
+    fail: common.mustCall((reason) => assert.strictEqual(reason, error)),
+  }), error);
+
+  await assert.rejects(pipeTo({
+    [Symbol.asyncIterator]() {
+      return { next: async () => 42 };
+    },
+  }, {
+    write: common.mustNotCall(),
+    writeSync: common.mustNotCall(),
+    fail: common.mustCall(),
+  }), { code: 'ERR_INVALID_RETURN_VALUE' });
+}
+
 Promise.all([
   testPipeToSync(),
   testPipeTo(),
@@ -412,5 +469,6 @@ Promise.all([
   testPipeToSyncIterableWriteFallback(),
   testPipeToSyncIterableAsyncValue(),
   testPipeToSyncIterableWriteError(),
+  testPipeToAsyncIterableErrors(),
   testPipeToSourceNormalizationIndependentOfWriter(),
 ]).then(common.mustCall());