stream: apply more stream/iter conformance fixes

Author: jasnell

lib/internal/streams/iter/consumers.js

@@ -10,16 +10,13 @@
 const {
   ArrayBufferPrototypeGetByteLength,
   ArrayBufferPrototypeSlice,
-  ArrayPrototypeFilter,
   ArrayPrototypeMap,
   ArrayPrototypePush,
   ArrayPrototypeSlice,
+  Promise,
   PromisePrototypeThen,
-  PromiseResolve,
   SafePromiseAllReturnVoid,
-  SafePromiseRace,
   SymbolAsyncIterator,
-  SymbolIterator,
   TypedArrayPrototypeGetBuffer,
   TypedArrayPrototypeGetByteLength,
   TypedArrayPrototypeGetByteOffset,
@@ -372,11 +369,7 @@ function ondrain(drainable) {
     return null;
   }
 
-  try {
-    return drainable[drainableProtocol]();
-  } catch {
-    return null;
-  }
+  return drainable[drainableProtocol]();
 }
 
 // =============================================================================
@@ -403,87 +396,105 @@ function merge(...args) {
     validateAbortSignal(options.signal, 'options.signal');
   }
 
+  // Normalize each source via from()
+  const normalized = ArrayPrototypeMap(sources, (source) => from(source));
+
   return {
     __proto__: null,
     async *[SymbolAsyncIterator]() {
       const signal = options?.signal;
 
       signal?.throwIfAborted();
 
-      if (sources.length === 0) return;
+      if (normalized.length === 0) return;
 
-      if (sources.length === 1) {
-        for await (const batch of sources[0]) {
+      if (normalized.length === 1) {
+        for await (const batch of normalized[0]) {
           signal?.throwIfAborted();
           yield batch;
         }
         return;
       }
 
-      // Multiple sources - race them
-      const states = ArrayPrototypeMap(sources, (source) => {
-        let iterator;
-        if (source[SymbolAsyncIterator]) {
-          iterator = source[SymbolAsyncIterator]();
-        } else if (source[SymbolIterator]) {
-          // Wrap sync iterator to async
-          const syncIter = source[SymbolIterator]();
-          iterator = {
-            __proto__: null,
-            next() { return PromiseResolve(syncIter.next()); },
-            return() {
-              return PromiseResolve(syncIter.return?.() ??
-                { __proto__: null, done: true, value: undefined });
-            },
-          };
+      // Multiple sources - use a ready queue so that batches that settle
+      // between consumer pulls are drained synchronously without an extra
+      // async tick per batch. Each source has at most one pending .next()
+      // at a time. Every batch from every source is preserved.
+      const ready = [];
+      let activeCount = normalized.length;
+      let waitResolve = null;
+
+      // Called when a source's .next() settles. Pushes the result into
+      // the ready queue and wakes the consumer if it's waiting.
+      const onSettled = (iterator, result) => {
+        if (result.done) {
+          activeCount--;
         } else {
-          throw new ERR_INVALID_ARG_TYPE(
-            'source', ['AsyncIterable', 'Iterable'], source);
+          ArrayPrototypePush(ready, result.value);
+          // Immediately request the next value from this source
+          // (at most one pending .next() per source)
+          PromisePrototypeThen(
+            iterator.next(),
+            (r) => onSettled(iterator, r),
+            (err) => {
+              ArrayPrototypePush(ready, { __proto__: null, error: err });
+              if (waitResolve) {
+                waitResolve();
+                waitResolve = null;
+              }
+            },
+          );
         }
-        return { __proto__: null, iterator, done: false, pending: null };
-      });
-
-      const startIterator = (state, index) => {
-        if (!state.done && !state.pending) {
-          state.pending = PromisePrototypeThen(
-            state.iterator.next(),
-            (result) => ({ __proto__: null, index, result }));
+        if (waitResolve) {
+          waitResolve();
+          waitResolve = null;
         }
       };
 
-      // Start all
-      for (let i = 0; i < states.length; i++) {
-        startIterator(states[i], i);
+      // Start one .next() per source
+      const iterators = [];
+      for (let i = 0; i < normalized.length; i++) {
+        const iterator = normalized[i][SymbolAsyncIterator]();
+        ArrayPrototypePush(iterators, iterator);
+        PromisePrototypeThen(
+          iterator.next(),
+          (r) => onSettled(iterator, r),
+          (err) => {
+            ArrayPrototypePush(ready, { __proto__: null, error: err });
+            if (waitResolve) {
+              waitResolve();
+              waitResolve = null;
+            }
+          },
+        );
       }
 
       try {
-        while (true) {
+        while (activeCount > 0 || ready.length > 0) {
           signal?.throwIfAborted();
 
-          const pending = ArrayPrototypeFilter(
-            ArrayPrototypeMap(states,
-                              (state) => state.pending),
-            (p) => p !== null);
-
-          if (pending.length === 0) break;
-
-          const { index, result } = await SafePromiseRace(pending);
-
-          states[index].pending = null;
+          // Drain ready queue synchronously
+          while (ready.length > 0) {
+            const item = ready.shift();
+            if (item?.error) {
+              throw item.error;
+            }
+            yield item;
+          }
 
-          if (result.done) {
-            states[index].done = true;
-          } else {
-            yield result.value;
-            startIterator(states[index], index);
+          // If sources are still active, wait for the next settlement
+          if (activeCount > 0) {
+            await new Promise((resolve) => {
+              waitResolve = resolve;
+            });
           }
         }
       } finally {
         // Clean up: return all iterators
-        await SafePromiseAllReturnVoid(states, async (state) => {
-          if (!state.done && state.iterator.return) {
+        await SafePromiseAllReturnVoid(iterators, async (iterator) => {
+          if (iterator.return) {
             try {
-              await state.iterator.return();
+              await iterator.return();
             } catch {
               // Ignore return errors
             }

test/parallel/test-stream-iter-consumers-merge.js

@@ -147,6 +147,21 @@ async function testMergeSignalMidIteration() {
   await assert.rejects(() => iter.next(), { name: 'AbortError' });
 }
 
+// merge() accepts string sources (normalized via from())
+async function testMergeStringSources() {
+  const batches = [];
+  for await (const batch of merge('hello', 'world')) {
+    batches.push(batch);
+  }
+  // Each string becomes a single-batch source
+  assert.strictEqual(batches.length >= 2, true);
+  const combined = new TextDecoder().decode(
+    Buffer.concat(batches.flat()));
+  // Both strings should appear (order may vary)
+  assert.ok(combined.includes('hello'));
+  assert.ok(combined.includes('world'));
+}
+
 Promise.all([
   testMergeTwoSources(),
   testMergeSingleSource(),
@@ -156,4 +171,5 @@ Promise.all([
   testMergeSourceError(),
   testMergeConsumerBreak(),
   testMergeSignalMidIteration(),
+  testMergeStringSources(),
 ]).then(common.mustCall());

test/parallel/test-stream-iter-push-writer.js

@@ -26,6 +26,18 @@ async function testOndrainNonDrainable() {
   assert.strictEqual(ondrain('string'), null);
 }
 
+async function testOndrainProtocolErrorPropagates() {
+  const badDrainable = {
+    [Symbol.for('Stream.drainableProtocol')]() {
+      throw new Error('protocol error');
+    },
+  };
+  assert.throws(
+    () => ondrain(badDrainable),
+    { message: 'protocol error' },
+  );
+}
+
 async function testWriteWithSignalRejects() {
   const { writer, readable } = push({ highWaterMark: 1 });
 
@@ -331,6 +343,7 @@ Promise.all([
   testWritevMixedTypes(),
   testWriteAfterEnd(),
   testWriteAfterFail(),
+  testOndrainProtocolErrorPropagates(),
   testFail(),
   testEndAsyncReturnValue(),
   testWriteUint8Array(),