From e027154707a253e155ac51d8dac5cd542c89c932 Mon Sep 17 00:00:00 2001
From: Nitzan Uziely <nitzan@testim.io>
Date: Sat, 13 Feb 2021 14:05:04 +0200
Subject: [PATCH] stream: add AbortSignal support to finished

Add AbortSignal support to stream.finished
---
 doc/api/stream.md                     |  6 +++
 lib/internal/streams/end-of-stream.js | 23 +++++++-
 test/parallel/test-stream-finished.js | 77 +++++++++++++++++++++++++++
 3 files changed, 105 insertions(+), 1 deletion(-)

diff --git a/doc/api/stream.md b/doc/api/stream.md
index cb4085806bb6da..80fdbb89e33ecd 100644
--- a/doc/api/stream.md
+++ b/doc/api/stream.md
@@ -1578,6 +1578,9 @@ further errors except from `_destroy()` may be emitted as `'error'`.
 <!-- YAML
 added: v10.0.0
 changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/37354
+    description: The `signal` option was added.
   - version: v14.0.0
     pr-url: https://github.com/nodejs/node/pull/32158
     description: The `finished(stream, cb)` will wait for the `'close'` event
@@ -1604,6 +1607,9 @@ changes:
   * `writable` {boolean} When set to `false`, the callback will be called when
     the stream ends even though the stream might still be writable.
     **Default:** `true`.
+  * `signal` {AbortSignal} allows aborting the wait for the stream finish. The
+    underlying stream will *not* be aborted if the signal is aborted. The
+    callback will get called with an `AbortError`.
 * `callback` {Function} A callback function that takes an optional error
   argument.
 * Returns: {Function} A cleanup function which removes all registered
diff --git a/lib/internal/streams/end-of-stream.js b/lib/internal/streams/end-of-stream.js
index 2f0043bf0b751e..96d5ba756e22ea 100644
--- a/lib/internal/streams/end-of-stream.js
+++ b/lib/internal/streams/end-of-stream.js
@@ -6,12 +6,18 @@
 const {
   FunctionPrototype,
   FunctionPrototypeCall,
+  ReflectApply,
 } = primordials;
+const {
+  AbortError,
+  codes,
+} = require('internal/errors');
 const {
   ERR_STREAM_PREMATURE_CLOSE
-} = require('internal/errors').codes;
+} = codes;
 const { once } = require('internal/util');
 const {
+  validateAbortSignal,
   validateFunction,
   validateObject,
 } = require('internal/validators');
@@ -76,6 +82,7 @@ function eos(stream, options, callback) {
     validateObject(options, 'options');
   }
   validateFunction(callback, 'callback');
+  validateAbortSignal(options.signal, 'options.signal');
 
   callback = once(callback);
 
@@ -199,6 +206,20 @@ function eos(stream, options, callback) {
     });
   }
 
+  if (options.signal && !closed) {
+    const abort = () => callback(new AbortError());
+    if (options.signal.aborted) {
+      process.nextTick(abort);
+    } else {
+      const originalCallback = callback;
+      callback = once((...args) => {
+        options.signal.removeEventListener('abort', abort);
+        ReflectApply(originalCallback, null, args);
+      });
+      options.signal.addEventListener('abort', abort);
+    }
+  }
+
   return function() {
     callback = nop;
     stream.removeListener('aborted', onclose);
diff --git a/test/parallel/test-stream-finished.js b/test/parallel/test-stream-finished.js
index a495bbe2861170..d1652396dae25d 100644
--- a/test/parallel/test-stream-finished.js
+++ b/test/parallel/test-stream-finished.js
@@ -92,6 +92,83 @@ const http = require('http');
   run();
 }
 
+{
+  // Check pre-cancelled
+  const signal = new EventTarget();
+  signal.aborted = true;
+
+  const rs = Readable.from((function* () {})());
+  finished(rs, { signal }, common.mustCall((err) => {
+    assert.strictEqual(err.name, 'AbortError');
+  }));
+}
+
+{
+  // Check cancelled before the stream ends sync.
+  const ac = new AbortController();
+  const { signal } = ac;
+
+  const rs = Readable.from((function* () {})());
+  finished(rs, { signal }, common.mustCall((err) => {
+    assert.strictEqual(err.name, 'AbortError');
+  }));
+
+  ac.abort();
+}
+
+{
+  // Check cancelled before the stream ends async.
+  const ac = new AbortController();
+  const { signal } = ac;
+
+  const rs = Readable.from((function* () {})());
+  setTimeout(() => ac.abort(), 1);
+  finished(rs, { signal }, common.mustCall((err) => {
+    assert.strictEqual(err.name, 'AbortError');
+  }));
+}
+
+{
+  // Check cancelled after doesn't throw.
+  const ac = new AbortController();
+  const { signal } = ac;
+
+  const rs = Readable.from((function* () {
+    yield 5;
+    setImmediate(() => ac.abort());
+  })());
+  rs.resume();
+  finished(rs, { signal }, common.mustSucceed());
+}
+
+{
+  // Promisified abort works
+  const finishedPromise = promisify(finished);
+  async function run() {
+    const ac = new AbortController();
+    const { signal } = ac;
+    const rs = Readable.from((function* () {})());
+    setImmediate(() => ac.abort());
+    await finishedPromise(rs, { signal });
+  }
+
+  assert.rejects(run, { name: 'AbortError' }).then(common.mustCall());
+}
+
+{
+  // Promisified pre-aborted works
+  const finishedPromise = promisify(finished);
+  async function run() {
+    const signal = new EventTarget();
+    signal.aborted = true;
+    const rs = Readable.from((function* () {})());
+    await finishedPromise(rs, { signal });
+  }
+
+  assert.rejects(run, { name: 'AbortError' }).then(common.mustCall());
+}
+
+
 {
   const rs = fs.createReadStream('file-does-not-exist');