From d06ab236dc1f572856996ed73202aca34350c973 Mon Sep 17 00:00:00 2001
From: Benjamin Gruenbaum <benjamingr@gmail.com>
Date: Sun, 16 Jun 2024 02:41:59 +0300
Subject: [PATCH] stream: support dispose in writable

Add support to Symbol.asyncDispose in writable streams.
Additionally add a test for writable, transform and duplex streams
who inherit from readable/writable to avoid breakage.

Co-authored-by: Robert Nagy <ronagy@icloud.com>
Co-authored-by: atlowChemi <chemi@atlow.co.il>
PR-URL: https://github.com/nodejs/node/pull/48547
Reviewed-By: Chemi Atlow <chemi@atlow.co.il>
Reviewed-By: Moshe Atlow <moshe@atlow.co.il>
Reviewed-By: Robert Nagy <ronagy@icloud.com>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
---
 doc/api/stream.md                             | 11 ++++++
 lib/internal/streams/writable.js              | 37 ++++++++++++++-----
 test/parallel/test-stream-duplex-destroy.js   | 15 ++++++++
 .../parallel/test-stream-transform-destroy.js | 11 ++++++
 test/parallel/test-stream-writable-destroy.js | 12 ++++++
 5 files changed, 76 insertions(+), 10 deletions(-)

diff --git a/doc/api/stream.md b/doc/api/stream.md
index f7566858ee594b..9539457daf8749 100644
--- a/doc/api/stream.md
+++ b/doc/api/stream.md
@@ -954,6 +954,17 @@ added: v12.3.0
 
 Getter for the property `objectMode` of a given `Writable` stream.
 
+##### `writable[Symbol.asyncDispose]()`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+Calls [`writable.destroy()`][writable-destroy] with an `AbortError` and returns
+a promise that fulfills when the stream is finished.
+
 ##### `writable.write(chunk[, encoding][, callback])`
 
 <!-- YAML
diff --git a/lib/internal/streams/writable.js b/lib/internal/streams/writable.js
index 268b54b964a835..96443f1c7898fb 100644
--- a/lib/internal/streams/writable.js
+++ b/lib/internal/streams/writable.js
@@ -32,8 +32,10 @@ const {
   ObjectDefineProperties,
   ObjectDefineProperty,
   ObjectSetPrototypeOf,
+  Promise,
   StringPrototypeToLowerCase,
   Symbol,
+  SymbolAsyncDispose,
   SymbolHasInstance,
 } = primordials;
 
@@ -44,6 +46,7 @@ const EE = require('events');
 const Stream = require('internal/streams/legacy').Stream;
 const { Buffer } = require('buffer');
 const destroyImpl = require('internal/streams/destroy');
+const eos = require('internal/streams/end-of-stream');
 
 const {
   addAbortSignal,
@@ -54,16 +57,19 @@ const {
   getDefaultHighWaterMark,
 } = require('internal/streams/state');
 const {
-  ERR_INVALID_ARG_TYPE,
-  ERR_METHOD_NOT_IMPLEMENTED,
-  ERR_MULTIPLE_CALLBACK,
-  ERR_STREAM_CANNOT_PIPE,
-  ERR_STREAM_DESTROYED,
-  ERR_STREAM_ALREADY_FINISHED,
-  ERR_STREAM_NULL_VALUES,
-  ERR_STREAM_WRITE_AFTER_END,
-  ERR_UNKNOWN_ENCODING,
-} = require('internal/errors').codes;
+  AbortError,
+  codes: {
+    ERR_INVALID_ARG_TYPE,
+    ERR_METHOD_NOT_IMPLEMENTED,
+    ERR_MULTIPLE_CALLBACK,
+    ERR_STREAM_ALREADY_FINISHED,
+    ERR_STREAM_CANNOT_PIPE,
+    ERR_STREAM_DESTROYED,
+    ERR_STREAM_NULL_VALUES,
+    ERR_STREAM_WRITE_AFTER_END,
+    ERR_UNKNOWN_ENCODING,
+  },
+} = require('internal/errors');
 const {
   kState,
   // bitfields
@@ -1142,3 +1148,14 @@ Writable.fromWeb = function(writableStream, options) {
 Writable.toWeb = function(streamWritable) {
   return lazyWebStreams().newWritableStreamFromStreamWritable(streamWritable);
 };
+
+Writable.prototype[SymbolAsyncDispose] = function() {
+  let error;
+  if (!this.destroyed) {
+    error = this.writableFinished ? null : new AbortError();
+    this.destroy(error);
+  }
+  return new Promise((resolve, reject) =>
+    eos(this, (err) => (err && err.name !== 'AbortError' ? reject(err) : resolve(null))),
+  );
+};
diff --git a/test/parallel/test-stream-duplex-destroy.js b/test/parallel/test-stream-duplex-destroy.js
index 9258cacd0d6f69..52867384057f59 100644
--- a/test/parallel/test-stream-duplex-destroy.js
+++ b/test/parallel/test-stream-duplex-destroy.js
@@ -269,3 +269,18 @@ const assert = require('assert');
   }));
   duplex.destroy();
 }
+
+{
+  // Check Symbol.asyncDispose
+  const duplex = new Duplex({
+    write(chunk, enc, cb) { cb(); },
+    read() {},
+  });
+  let count = 0;
+  duplex.on('error', common.mustCall((e) => {
+    assert.strictEqual(count++, 0); // Ensure not called twice
+    assert.strictEqual(e.name, 'AbortError');
+  }));
+  duplex.on('close', common.mustCall());
+  duplex[Symbol.asyncDispose]().then(common.mustCall());
+}
diff --git a/test/parallel/test-stream-transform-destroy.js b/test/parallel/test-stream-transform-destroy.js
index 8750c65164e5e0..428bab9ce33fcf 100644
--- a/test/parallel/test-stream-transform-destroy.js
+++ b/test/parallel/test-stream-transform-destroy.js
@@ -141,3 +141,14 @@ const assert = require('assert');
 
   transform.destroy();
 }
+
+{
+  const transform = new Transform({
+    transform(chunk, enc, cb) {}
+  });
+  transform.on('error', common.mustCall((err) => {
+    assert.strictEqual(err.name, 'AbortError');
+  }));
+  transform.on('close', common.mustCall());
+  transform[Symbol.asyncDispose]().then(common.mustCall());
+}
diff --git a/test/parallel/test-stream-writable-destroy.js b/test/parallel/test-stream-writable-destroy.js
index 768ad212b37e89..05d7932b88c182 100644
--- a/test/parallel/test-stream-writable-destroy.js
+++ b/test/parallel/test-stream-writable-destroy.js
@@ -487,3 +487,15 @@ const assert = require('assert');
   }));
   s.destroy(_err);
 }
+
+{
+  const write = new Writable({
+    write(chunk, enc, cb) { cb(); }
+  });
+
+  write.on('error', common.mustCall((e) => {
+    assert.strictEqual(e.name, 'AbortError');
+    assert.strictEqual(write.destroyed, true);
+  }));
+  write[Symbol.asyncDispose]().then(common.mustCall());
+}