stream: optimize forEach(), some(), every() and find() functions on Readable

lib/internal/streams/operators.js

@@ -3,20 +3,23 @@
 const {
   ArrayPrototypePush,
   Boolean,
+  FunctionPrototypeCall,
   MathFloor,
   Number,
   NumberIsNaN,
   Promise,
   PromisePrototypeThen,
   PromiseReject,
   PromiseResolve,
+  PromiseWithResolvers,
   Symbol,
 } = primordials;
 
 const { AbortController, AbortSignal } = require('internal/abort_controller');
 
 const {
   AbortError,
+  aggregateTwoErrors,
   codes: {
     ERR_MISSING_ARGS,
     ERR_OUT_OF_RANGE,
@@ -27,14 +30,44 @@ const {
   validateInteger,
   validateObject,
   validateFunction,
+  validateBoolean,
 } = require('internal/validators');
 const { kWeakHandler, kResistStopPropagation } = require('internal/event_target');
 const destroyImpl = require('internal/streams/destroy');
-const { finished } = require('internal/streams/end-of-stream');
+const { eos, finished } = require('internal/streams/end-of-stream');
 
 const kEmpty = Symbol('kEmpty');
 const kEof = Symbol('kEof');
 
+const {
+  isPromise,
+} = require('internal/util/types');
+
+
+// Return native promises unchanged and normalize custom thenables exactly once.
+// Capturing `then` avoids repeated getter access, while calling it with `value`
+// as the receiver preserves thenables that depend on their `this` value.
+function getThenablePromise(value) {
+  if (isPromise(value)) {
+    return value;
+  }
+
+  const valueType = typeof value;
+  if ((valueType === 'object' && value !== null) || valueType === 'function') {
+    const then = value.then;
+    if (typeof then === 'function') {
+      return PromiseResolve({
+        __proto__: null,
+        then(resolve, reject) {
+          FunctionPrototypeCall(then, value, resolve, reject);
+        },
+      });
+    }
+  }
+
+  return undefined;
+}
+
 function map(fn, options) {
   validateFunction(fn, 'fn');
   if (options != null) {
@@ -183,36 +216,208 @@ function map(fn, options) {
   }.call(this);
 }
 
-async function some(fn, options = undefined) {
-  for await (const unused of filter.call(this, fn, options)) {
-    return true;
+function nowOrLater(fn, fn2, args) {
+  const value = fn(...args);
+  const promise = getThenablePromise(value);
+  if (promise !== undefined) {
+    return PromisePrototypeThen(promise, fn2);
   }
-  return false;
+  return fn2(value);
+}
+
+async function some(fn, options = undefined) {
+  validateFunction(fn, 'fn');
+  const someFn = (...args) => {
+    return nowOrLater(fn, Boolean, args);
+  };
+  return (await find.call(this, someFn, options)) !== undefined;
 }
 
 async function every(fn, options = undefined) {
   validateFunction(fn, 'fn');
+  const everyFn = (...args) => {
+    return nowOrLater(fn, (value) => !value, args);
+  };
   // https://en.wikipedia.org/wiki/De_Morgan%27s_laws
-  return !(await some.call(this, async (...args) => {
-    return !(await fn(...args));
-  }, options));
+  return !(await find.call(this, everyFn, options));
 }
 
-async function find(fn, options) {
-  for await (const result of filter.call(this, fn, options)) {
-    return result;
+function find(fn, options) {
+  validateFunction(fn, 'fn');
+
+  if (options != null) {
+    validateObject(options, 'options');
   }
-  return undefined;
+  const signal = options?.signal;
+  if (signal != null) {
+    validateAbortSignal(signal, 'options.signal');
+  }
+
+  const concurrency = MathFloor(options?.concurrency ?? 1);
+  validateInteger(concurrency, 'options.concurrency', 1);
+
+  const destroyOnReturn = options?.destroyOnReturn ?? true;
+  validateBoolean(destroyOnReturn, 'options.destroyOnReturn');
+
+  const ac = new AbortController();
+  const predicateSignal = AbortSignal.any([ac.signal, signal].filter(Boolean));
+  const predicateOptions = { signal: predicateSignal };
+
+  // Concurrent predicates can settle out of order. Stop reading after any
+  // match, but keep the lowest index after all active predicates settle.
+  const stream = this;
+  const { promise, resolve } = PromiseWithResolvers();
+  let match;
+  let error;
+  let activeEvaluations = 0;
+  let nextIndex = 0;
+  let ended = false;
+  let settled = false;
+  let draining = false;
+
+  function settle() {
+    if (!settled) {
+      settled = true;
+      resolve();
+    }
+  }
+
+  function fail(err) {
+    if (settled) {
+      return;
+    }
+    error = aggregateTwoErrors(error, err);
+    destroyImpl.destroyer(stream, error);
+    settle();
+  }
+
+  function maybeSettle() {
+    if (activeEvaluations === 0 && (match !== undefined || ended)) {
+      settle();
+    }
+  }
+
+  function evaluationFinished(matches, chunk, index) {
+    if (matches && (match === undefined || index < match.index)) {
+      match = { index, value: chunk };
+    }
+    activeEvaluations--;
+    maybeSettle();
+
+    if (!settled && match === undefined && !draining) {
+      onReadable();
+    }
+  }
+
+  function evaluationRejected(err) {
+    activeEvaluations--;
+    fail(err);
+  }
+
+  function evaluate(chunk, index) {
+    activeEvaluations++;
+
+    let matches;
+    try {
+      matches = fn(chunk, predicateOptions);
+      const matchesPromise = getThenablePromise(matches);
+      if (matchesPromise !== undefined) {
+        PromisePrototypeThen(
+          matchesPromise,
+          (result) => evaluationFinished(result, chunk, index),
+          evaluationRejected,
+        );
+        return;
+      }
+    } catch (err) {
+      evaluationRejected(err);
+      return;
+    }
+    evaluationFinished(matches, chunk, index);
+  }
+
+  function onReadable() {
+    if (draining || settled || match !== undefined) {
+      return;
+    }
+
+    draining = true;
+    try {
+      while (
+        !settled &&
+        match === undefined &&
+        activeEvaluations < concurrency
+      ) {
+        if (signal?.aborted) {
+          fail(new AbortError(undefined, { cause: signal.reason }));
+          return;
+        }
+
+        const chunk = stream.destroyed ? null : stream.read();
+        if (chunk === null) {
+          return;
+        }
+        evaluate(chunk, nextIndex++);
+      }
+    } catch (err) {
+      fail(err);
+    } finally {
+      draining = false;
+    }
+  }
+
+  function onAbort() {
+    fail(new AbortError(undefined, { cause: signal.reason }));
+  }
+
+  stream.on('readable', onReadable);
+
+  const cleanup = eos(stream, { writable: false }, (err) => {
+    if (settled) {
+      return;
+    }
+    if (err) {
+      fail(err);
+      return;
+    }
+    ended = true;
+    maybeSettle();
+  });
+
+  if (signal != null) {
+    signal.addEventListener('abort', onAbort, { once: true });
+    if (signal.aborted) {
+      onAbort();
+    }
+  }
+
+  return PromisePrototypeThen(promise, () => {
+    stream.off('readable', onReadable);
+    signal?.removeEventListener('abort', onAbort);
+    ac.abort();
+
+    if (
+      (error || destroyOnReturn !== false) &&
+      (error === undefined || stream._readableState.autoDestroy)
+    ) {
+      destroyImpl.destroyer(stream, error);
+    } else {
+      cleanup();
+    }
+
+    if (error) {
+      return PromiseReject(error);
+    }
+    return match?.value;
+  });
 }
 
 async function forEach(fn, options) {
   validateFunction(fn, 'fn');
-  async function forEachFn(value, options) {
-    await fn(value, options);
-    return kEmpty;
-  }
-  // eslint-disable-next-line no-unused-vars
-  for await (const unused of map.call(this, forEachFn, options));
+  const forEachFn = (...args) => {
+    return nowOrLater(fn, () => false, args);
+  };
+  await find.call(this, forEachFn, options);
 }
 
 function filter(fn, options) {

test/parallel/test-stream-forEach.js

@@ -43,6 +43,30 @@ const { once } = require('events');
   })().then(common.mustCall());
 }
 
+{
+  // forEach awaits thenables returned by the callback.
+  const visited = [];
+  const receivers = [];
+  const thenables = [];
+  (async () => {
+    await Readable.from([1, 2]).forEach((value) => {
+      const thenable = {
+        then(resolve) {
+          receivers.push(this);
+          setImmediate(() => {
+            visited.push(value);
+            resolve();
+          });
+        },
+      };
+      thenables.push(thenable);
+      return thenable;
+    });
+    assert.deepStrictEqual(visited, [1, 2]);
+    assert.deepStrictEqual(receivers, thenables);
+  })().then(common.mustCall());
+}
+
 {
   // forEach works on an infinite stream
   const ac = new AbortController();