From dc74b355953f1518793784c26ae318d88ab783ca Mon Sep 17 00:00:00 2001
From: Nitzan Uziely <nitzan@testim.io>
Date: Sat, 13 Feb 2021 17:15:23 +0200
Subject: [PATCH] stream: add AbortSignal to promisified pipeline

add support for AbortSignal to promisified pipeline.

Resolves: https://github.com/nodejs/node/issues/37321
---
 doc/api/stream.md                     | 30 ++++++++++-
 lib/stream/promises.js                | 45 ++++++++++++++++-
 test/parallel/test-stream-pipeline.js | 72 +++++++++++++++++++++++++++
 3 files changed, 145 insertions(+), 2 deletions(-)

diff --git a/doc/api/stream.md b/doc/api/stream.md
index 1d6b811c655fd5..3818913ada8cd0 100644
--- a/doc/api/stream.md
+++ b/doc/api/stream.md
@@ -1719,7 +1719,11 @@ pipeline(
 );
 ```
 
-The `pipeline` API provides promise version:
+The `pipeline` API provides a promise version, which can also
+receive an options argument as the last parameter with a
+`signal` {AbortSignal} property. When the signal is aborted,
+`destroy` will be called on the underlying pipeline, with an
+`AbortError`.
 
 ```js
 const { pipeline } = require('stream/promises');
@@ -1736,6 +1740,30 @@ async function run() {
 run().catch(console.error);
 ```
 
+To use an `AbortSignal`, pass it inside an options object,
+as the last argument:
+
+```js
+const { pipeline } = require('stream/promises');
+
+async function run() {
+  const ac = new AbortController();
+  const options = {
+    signal: ac.signal,
+  };
+
+  setTimeout(() => ac.abort(), 1);
+  await pipeline(
+    fs.createReadStream('archive.tar'),
+    zlib.createGzip(),
+    fs.createWriteStream('archive.tar.gz'),
+    options,
+  );
+}
+
+run().catch(console.error); // AbortError
+```
+
 The `pipeline` API also supports async generators:
 
 ```js
diff --git a/lib/stream/promises.js b/lib/stream/promises.js
index 986db2e1f8db8a..90d6b1244bd380 100644
--- a/lib/stream/promises.js
+++ b/lib/stream/promises.js
@@ -1,22 +1,65 @@
 'use strict';
 
 const {
+  ArrayPrototypePop,
   Promise,
+  SymbolAsyncIterator,
+  SymbolIterator,
 } = primordials;
 
+const {
+  addAbortSignalNoValidate,
+} = require('internal/streams/add-abort-signal');
+
+const {
+  validateAbortSignal,
+} = require('internal/validators');
+
 let pl;
 let eos;
 
+function isReadable(obj) {
+  return !!(obj && typeof obj.pipe === 'function');
+}
+
+function isWritable(obj) {
+  return !!(obj && typeof obj.write === 'function');
+}
+
+function isStream(obj) {
+  return isReadable(obj) || isWritable(obj);
+}
+
+function isIterable(obj, isAsync) {
+  if (!obj) return false;
+  if (isAsync === true) return typeof obj[SymbolAsyncIterator] === 'function';
+  if (isAsync === false) return typeof obj[SymbolIterator] === 'function';
+  return typeof obj[SymbolAsyncIterator] === 'function' ||
+    typeof obj[SymbolIterator] === 'function';
+}
+
 function pipeline(...streams) {
   if (!pl) pl = require('internal/streams/pipeline');
   return new Promise((resolve, reject) => {
-    pl(...streams, (err, value) => {
+    let signal;
+    const lastArg = streams[streams.length - 1];
+    if (lastArg && typeof lastArg === 'object' &&
+        !isStream(lastArg) && !isIterable(lastArg)) {
+      const options = ArrayPrototypePop(streams);
+      signal = options.signal;
+      validateAbortSignal(signal, 'options.signal');
+    }
+
+    const pipe = pl(...streams, (err, value) => {
       if (err) {
         reject(err);
       } else {
         resolve(value);
       }
     });
+    if (signal) {
+      addAbortSignalNoValidate(signal, pipe);
+    }
   });
 }
 
diff --git a/test/parallel/test-stream-pipeline.js b/test/parallel/test-stream-pipeline.js
index 78057f9eeffec6..86f9006d83c430 100644
--- a/test/parallel/test-stream-pipeline.js
+++ b/test/parallel/test-stream-pipeline.js
@@ -469,6 +469,78 @@ const net = require('net');
   run();
 }
 
+{
+  // Check aborted signal without values
+  const pipelinePromise = promisify(pipeline);
+  async function run() {
+    const ac = new AbortController();
+    const { signal } = ac;
+    async function* producer() {
+      ac.abort();
+      await Promise.resolve();
+      yield '8';
+    }
+
+    const w = new Writable({
+      write(chunk, encoding, callback) {
+        callback();
+      }
+    });
+    await pipelinePromise(producer, w, { signal });
+  }
+
+  assert.rejects(run, { name: 'AbortError' }).then(common.mustCall());
+}
+
+{
+  // Check aborted signal after init.
+  const pipelinePromise = promisify(pipeline);
+  async function run() {
+    const ac = new AbortController();
+    const { signal } = ac;
+    async function* producer() {
+      yield '5';
+      await Promise.resolve();
+      ac.abort();
+      await Promise.resolve();
+      yield '8';
+    }
+
+    const w = new Writable({
+      write(chunk, encoding, callback) {
+        callback();
+      }
+    });
+    await pipelinePromise(producer, w, { signal });
+  }
+
+  assert.rejects(run, { name: 'AbortError' }).then(common.mustCall());
+}
+
+{
+  // Check pre-aborted signal
+  const pipelinePromise = promisify(pipeline);
+  async function run() {
+    const ac = new AbortController();
+    const { signal } = ac;
+    ac.abort();
+    async function* producer() {
+      yield '5';
+      await Promise.resolve();
+      yield '8';
+    }
+
+    const w = new Writable({
+      write(chunk, encoding, callback) {
+        callback();
+      }
+    });
+    await pipelinePromise(producer, w, { signal });
+  }
+
+  assert.rejects(run, { name: 'AbortError' }).then(common.mustCall());
+}
+
 {
   const read = new Readable({
     read() {}