From e0807fd3718bc7b559c2c8e92588e240ba0a58b7 Mon Sep 17 00:00:00 2001
From: Darshan Sen <raisinten@gmail.com>
Date: Mon, 2 May 2022 10:09:39 +0530
Subject: [PATCH] worker: add hasRef() to MessagePort

Since we were removing the hasRef() method before exposing the
MessagePort object, the only way of knowing if the handle was keeping
the event loop active was to parse the string returned by
util.inspect(port), which is inconvenient and inconsistent with most of
the other async resources. So this change stops removing hasRef() from
the MessagePort prototype. The reason why this is also being documented
is that while reporting active resources, async_hooks returns the same
MessagePort object as the one that is accessible by users.

Refs: https://github.com/nodejs/node/issues/42091#issuecomment-1104793189
Signed-off-by: Darshan Sen <raisinten@gmail.com>

PR-URL: https://github.com/nodejs/node/pull/42849
Reviewed-By: Anna Henningsen <anna@addaleax.net>
---
 doc/api/worker_threads.md                     | 12 ++++
 lib/internal/worker/io.js                     |  5 +-
 test/parallel/test-messageport-hasref.js      | 57 +++++++++++++++++++
 test/parallel/test-worker-message-port.js     |  4 +-
 .../test-worker-messageport-hasref.js         | 45 +++++++++++++++
 5 files changed, 120 insertions(+), 3 deletions(-)
 create mode 100644 test/parallel/test-messageport-hasref.js
 create mode 100644 test/parallel/test-worker-messageport-hasref.js

diff --git a/doc/api/worker_threads.md b/doc/api/worker_threads.md
index 0fa741c6ed5325..0770edd28f3e41 100644
--- a/doc/api/worker_threads.md
+++ b/doc/api/worker_threads.md
@@ -757,6 +757,18 @@ port2.postMessage(new URL('https://example.org'));
 // Prints: { }
 ```
 
+### `port.hasRef()`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+* Returns: {boolean}
+
+If true, the `MessagePort` object will keep the Node.js event loop active.
+
 ### `port.ref()`
 
 <!-- YAML
diff --git a/lib/internal/worker/io.js b/lib/internal/worker/io.js
index a2d6fe42a32f1a..5b16c26443c466 100644
--- a/lib/internal/worker/io.js
+++ b/lib/internal/worker/io.js
@@ -91,7 +91,7 @@ const messageTypes = {
 // We have to mess with the MessagePort prototype a bit, so that a) we can make
 // it inherit from NodeEventTarget, even though it is a C++ class, and b) we do
 // not provide methods that are not present in the Browser and not documented
-// on our side (e.g. hasRef).
+// on our side (e.g. stopMessagePort).
 // Save a copy of the original set of methods as a shallow clone.
 const MessagePortPrototype = ObjectCreate(
   ObjectGetPrototypeOf(MessagePort.prototype),
@@ -103,6 +103,9 @@ ObjectSetPrototypeOf(MessagePort.prototype, NodeEventTarget.prototype);
 // changing the prototype of MessagePort.prototype implicitly removed them.
 MessagePort.prototype.ref = MessagePortPrototype.ref;
 MessagePort.prototype.unref = MessagePortPrototype.unref;
+MessagePort.prototype.hasRef = function() {
+  return !!FunctionPrototypeCall(MessagePortPrototype.hasRef, this);
+};
 
 function validateMessagePort(port, name) {
   if (!checkMessagePort(port))
diff --git a/test/parallel/test-messageport-hasref.js b/test/parallel/test-messageport-hasref.js
new file mode 100644
index 00000000000000..bc213f7897d277
--- /dev/null
+++ b/test/parallel/test-messageport-hasref.js
@@ -0,0 +1,57 @@
+'use strict';
+const common = require('../common');
+
+const { MessageChannel } = require('worker_threads');
+const { createHook } = require('async_hooks');
+const { strictEqual } = require('assert');
+
+const handles = [];
+
+createHook({
+  init(asyncId, type, triggerAsyncId, resource) {
+    if (type === 'MESSAGEPORT') {
+      handles.push(resource);
+    }
+  }
+}).enable();
+
+const { port1, port2 } = new MessageChannel();
+strictEqual(handles[0], port1);
+strictEqual(handles[1], port2);
+
+strictEqual(handles[0].hasRef(), false);
+strictEqual(handles[1].hasRef(), false);
+
+port1.unref();
+strictEqual(handles[0].hasRef(), false);
+
+port1.ref();
+strictEqual(handles[0].hasRef(), true);
+
+port1.unref();
+strictEqual(handles[0].hasRef(), false);
+
+port1.on('message', () => {});
+strictEqual(handles[0].hasRef(), true);
+
+port2.unref();
+strictEqual(handles[1].hasRef(), false);
+
+port2.ref();
+strictEqual(handles[1].hasRef(), true);
+
+port2.unref();
+strictEqual(handles[1].hasRef(), false);
+
+port2.on('message', () => {});
+strictEqual(handles[0].hasRef(), true);
+
+port1.on('close', common.mustCall(() => {
+  strictEqual(handles[0].hasRef(), false);
+  strictEqual(handles[1].hasRef(), false);
+}));
+
+port2.close();
+
+strictEqual(handles[0].hasRef(), true);
+strictEqual(handles[1].hasRef(), true);
diff --git a/test/parallel/test-worker-message-port.js b/test/parallel/test-worker-message-port.js
index ca28f6ccb73247..2663dde2a1b75f 100644
--- a/test/parallel/test-worker-message-port.js
+++ b/test/parallel/test-worker-message-port.js
@@ -179,7 +179,7 @@ const { MessageChannel, MessagePort } = require('worker_threads');
   assert.deepStrictEqual(
     Object.getOwnPropertyNames(MessagePort.prototype).sort(),
     [
-      'close', 'constructor', 'onmessage', 'onmessageerror', 'postMessage',
-      'ref', 'start', 'unref',
+      'close', 'constructor', 'hasRef', 'onmessage', 'onmessageerror',
+      'postMessage', 'ref', 'start', 'unref',
     ]);
 }
diff --git a/test/parallel/test-worker-messageport-hasref.js b/test/parallel/test-worker-messageport-hasref.js
new file mode 100644
index 00000000000000..448787742e308d
--- /dev/null
+++ b/test/parallel/test-worker-messageport-hasref.js
@@ -0,0 +1,45 @@
+'use strict';
+const common = require('../common');
+
+const { Worker } = require('worker_threads');
+const { createHook } = require('async_hooks');
+const { deepStrictEqual, strictEqual } = require('assert');
+
+const m = new Map();
+createHook({
+  init(asyncId, type, triggerAsyncId, resource) {
+    if (['WORKER', 'MESSAGEPORT'].includes(type)) {
+      m.set(asyncId, { type, resource });
+    }
+  },
+  destroy(asyncId) {
+    m.delete(asyncId);
+  }
+}).enable();
+
+function getActiveWorkerAndMessagePortTypes() {
+  const activeWorkerAndMessagePortTypes = [];
+  for (const asyncId of m.keys()) {
+    const { type, resource } = m.get(asyncId);
+    // Same logic as https://github.com/mafintosh/why-is-node-running/blob/24fb4c878753390a05d00959e6173d0d3c31fddd/index.js#L31-L32.
+    if (typeof resource.hasRef !== 'function' || resource.hasRef() === true) {
+      activeWorkerAndMessagePortTypes.push(type);
+    }
+  }
+  return activeWorkerAndMessagePortTypes;
+}
+
+const w = new Worker('', { eval: true });
+deepStrictEqual(getActiveWorkerAndMessagePortTypes(), ['WORKER']);
+w.unref();
+deepStrictEqual(getActiveWorkerAndMessagePortTypes(), []);
+w.ref();
+deepStrictEqual(getActiveWorkerAndMessagePortTypes(), ['WORKER', 'MESSAGEPORT']);
+
+w.on('exit', common.mustCall((exitCode) => {
+  strictEqual(exitCode, 0);
+  deepStrictEqual(getActiveWorkerAndMessagePortTypes(), ['WORKER']);
+  setTimeout(common.mustCall(() => {
+    deepStrictEqual(getActiveWorkerAndMessagePortTypes(), []);
+  }), 0);
+}));