From 03ddd13d8a72190aa22194eb535e9cb683f4f815 Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna@addaleax.net>
Date: Wed, 7 Feb 2018 01:36:20 +0100
Subject: [PATCH] net: use `_final` instead of `on('finish')`

Shutting down the connection is what `_final` is there for.

PR-URL: https://github.com/nodejs/node/pull/18608
Reviewed-By: Luigi Pinca <luigipinca@gmail.com>
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
---
 lib/internal/streams/destroy.js               |  2 ++
 lib/net.js                                    | 28 +++++++++++--------
 test/async-hooks/test-shutdownwrap.js         | 12 ++++----
 test/parallel/test-stream-writable-destroy.js | 14 ++++++++++
 4 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/lib/internal/streams/destroy.js b/lib/internal/streams/destroy.js
index 86a22633f2fe78..985332ac4607a8 100644
--- a/lib/internal/streams/destroy.js
+++ b/lib/internal/streams/destroy.js
@@ -55,6 +55,8 @@ function undestroy() {
     this._writableState.destroyed = false;
     this._writableState.ended = false;
     this._writableState.ending = false;
+    this._writableState.finalCalled = false;
+    this._writableState.prefinished = false;
     this._writableState.finished = false;
     this._writableState.errorEmitted = false;
   }
diff --git a/lib/net.js b/lib/net.js
index ed2ea2dc27739b..edf8a6c0bc545b 100644
--- a/lib/net.js
+++ b/lib/net.js
@@ -256,7 +256,6 @@ function Socket(options) {
   }
 
   // shut down the socket when we're finished with it.
-  this.on('finish', onSocketFinish);
   this.on('_socketEnd', onSocketEnd);
 
   initSocketHandle(this);
@@ -303,39 +302,42 @@ Socket.prototype._unrefTimer = function _unrefTimer() {
 
 function shutdownSocket(self, callback) {
   var req = new ShutdownWrap();
-  req.oncomplete = callback;
+  req.oncomplete = afterShutdown;
   req.handle = self._handle;
+  req.callback = callback;
   return self._handle.shutdown(req);
 }
 
 // the user has called .end(), and all the bytes have been
 // sent out to the other side.
-function onSocketFinish() {
-  // If still connecting - defer handling 'finish' until 'connect' will happen
+Socket.prototype._final = function(cb) {
+  // If still connecting - defer handling `_final` until 'connect' will happen
   if (this.connecting) {
-    debug('osF: not yet connected');
-    return this.once('connect', onSocketFinish);
+    debug('_final: not yet connected');
+    return this.once('connect', () => this._final(cb));
   }
 
-  debug('onSocketFinish');
   if (!this.readable || this._readableState.ended) {
-    debug('oSF: ended, destroy', this._readableState);
+    debug('_final: ended, destroy', this._readableState);
+    cb();
     return this.destroy();
   }
 
-  debug('oSF: not ended, call shutdown()');
+  debug('_final: not ended, call shutdown()');
 
   // otherwise, just shutdown, or destroy() if not possible
-  if (!this._handle || !this._handle.shutdown)
+  if (!this._handle || !this._handle.shutdown) {
+    cb();
     return this.destroy();
+  }
 
   var err = defaultTriggerAsyncIdScope(
-    this[async_id_symbol], shutdownSocket, this, afterShutdown
+    this[async_id_symbol], shutdownSocket, this, cb
   );
 
   if (err)
     return this.destroy(errnoException(err, 'shutdown'));
-}
+};
 
 
 function afterShutdown(status, handle) {
@@ -344,6 +346,8 @@ function afterShutdown(status, handle) {
   debug('afterShutdown destroyed=%j', self.destroyed,
         self._readableState);
 
+  this.callback();
+
   // callback may come after call to destroy.
   if (self.destroyed)
     return;
diff --git a/test/async-hooks/test-shutdownwrap.js b/test/async-hooks/test-shutdownwrap.js
index dfaac2a1c05a70..fea4a3a166a270 100644
--- a/test/async-hooks/test-shutdownwrap.js
+++ b/test/async-hooks/test-shutdownwrap.js
@@ -24,11 +24,13 @@ let endedConnection = false;
 function onconnection(c) {
   assert.strictEqual(hooks.activitiesOfTypes('SHUTDOWNWRAP').length, 0);
   c.end();
-  endedConnection = true;
-  const as = hooks.activitiesOfTypes('SHUTDOWNWRAP');
-  assert.strictEqual(as.length, 1);
-  checkInvocations(as[0], { init: 1 }, 'after ending client connection');
-  this.close(onserverClosed);
+  process.nextTick(() => {
+    endedConnection = true;
+    const as = hooks.activitiesOfTypes('SHUTDOWNWRAP');
+    assert.strictEqual(as.length, 1);
+    checkInvocations(as[0], { init: 1 }, 'after ending client connection');
+    this.close(onserverClosed);
+  });
 }
 
 function onconnected() {
diff --git a/test/parallel/test-stream-writable-destroy.js b/test/parallel/test-stream-writable-destroy.js
index 87e55eccc3f6fd..46c48511177813 100644
--- a/test/parallel/test-stream-writable-destroy.js
+++ b/test/parallel/test-stream-writable-destroy.js
@@ -185,3 +185,17 @@ const { inherits } = require('util');
     assert.strictEqual(expected, err);
   }));
 }
+
+{
+  // Checks that `._undestroy()` restores the state so that `final` will be
+  // called again.
+  const write = new Writable({
+    write: common.mustNotCall(),
+    final: common.mustCall((cb) => cb(), 2)
+  });
+
+  write.end();
+  write.destroy();
+  write._undestroy();
+  write.end();
+}