MessagePort: release an explicit ref() when the last 'message' listener is removed

src/jsc/bindings/webcore/JSMessagePort.cpp

@@ -203,12 +203,12 @@ static inline bool setJSMessagePort_onmessageSetter(JSGlobalObject& lexicalGloba
     vm.writeBarrier(&thisObject, value);
     ensureStillAliveHere(value);
 
-    // node: a callable handler starts the port and keeps the loop alive; assigning anything else
-    // clears the handler and lets the loop exit again.
+    // node: a callable handler starts the port and keeps the loop alive. Assigning anything
+    // else just removes the handler; MessagePort::onDidChangeListenerImpl releases the refs
+    // only if that removed the last 'message' listener (so `onmessage = null` with nothing
+    // set, or with other listeners still attached, changes nothing, as in node).
     if (value.isCallable())
         thisObject.wrapped().jsRef(&lexicalGlobalObject);
-    else
-        thisObject.wrapped().jsUnref(&lexicalGlobalObject);
 
     return true;
 }
@@ -354,7 +354,7 @@ static inline JSC::EncodedJSValue jsMessagePortPrototypeFunction_closeBody(JSC::
     UNUSED_PARAM(throwScope);
     UNUSED_PARAM(callFrame);
     auto& impl = castedThis->wrapped();
-    impl.jsUnref(lexicalGlobalObject);
+    impl.jsUnref();
     RELEASE_AND_RETURN(throwScope, JSValue::encode(toJS<IDLUndefined>(*lexicalGlobalObject, throwScope, [&]() -> decltype(auto) { return impl.close(); })));
 }
 
@@ -385,7 +385,7 @@ static inline JSC::EncodedJSValue jsMessagePortPrototypeFunction_unrefBody(JSC::
     UNUSED_PARAM(throwScope);
     UNUSED_PARAM(callFrame);
     auto& impl = castedThis->wrapped();
-    RELEASE_AND_RETURN(throwScope, JSValue::encode(toJS<IDLUndefined>(*lexicalGlobalObject, throwScope, [&]() -> decltype(auto) { return impl.jsUnref(lexicalGlobalObject); })));
+    RELEASE_AND_RETURN(throwScope, JSValue::encode(toJS<IDLUndefined>(*lexicalGlobalObject, throwScope, [&]() -> decltype(auto) { return impl.jsUnref(); })));
 }
 
 JSC_DEFINE_HOST_FUNCTION(jsMessagePortPrototypeFunction_unref, (JSGlobalObject * lexicalGlobalObject, CallFrame* callFrame))

src/jsc/bindings/webcore/MessagePort.cpp

@@ -224,15 +224,8 @@ void MessagePort::close()
     // it in hasPendingActivity()); marking our side Closed is sufficient.
     m_pipe->close(m_side, MessagePortPipe::CloseKind::Explicit);
 
-    // Release the self-reference taken by jsRef() (set when .onmessage is
-    // assigned or .ref() is called from JS). The JS .close() binding calls
-    // jsUnref() first; stop() and contextDestroyed() do not.
-    if (m_hasRef) {
-        m_hasRef = false;
-        if (auto* context = scriptExecutionContext())
-            context->unrefEventLoop();
-        deref();
-    }
+    // The JS .close() binding calls jsUnref() first; stop() and contextDestroyed() do not.
+    releaseJsRef();
 
     // close() can run without a prior jsUnref() (warn-and-close, contextDestroyed());
     // clear the listener keepalive so a later listener add can't re-ref the loop.
@@ -287,8 +280,7 @@ void MessagePort::peerClosed()
     dispatchCloseEvent();
     // jsUnref() clears both the listener loop-ref (m_isRefd) and the onmessage/ref()
     // keepalive (m_hasRef), so a listening transferred port stops pinning the loop.
-    auto* globalObject = defaultGlobalObject(context->globalObject());
-    jsUnref(globalObject);
+    jsUnref();
 }
 
 TransferredMessagePort MessagePort::disentangle()
@@ -301,17 +293,11 @@ TransferredMessagePort MessagePort::disentangle()
     removeAllEventListeners();
     m_hasMessageEventListener = false;
 
-    // Release the self-reference taken by jsRef() on the sending side. After
-    // transfer this object is inert (the receiving side gets a fresh
+    // After transfer this object is inert (the receiving side gets a fresh
     // MessagePort for the same pipe endpoint) and is no longer a destruction
-    // observer, so nothing else will ever release a ref taken here.
-    // The caller (disentanglePorts) holds a RefPtr, so deref() is safe.
-    if (m_hasRef) {
-        m_hasRef = false;
-        if (auto* context = scriptExecutionContext())
-            context->unrefEventLoop();
-        deref();
-    }
+    // observer, so nothing else would ever release a jsRef() taken on it.
+    // The caller (disentanglePorts) holds a RefPtr, so the deref() is safe.
+    releaseJsRef();
 
     // A transferred port is inert; clear the listener keepalive too so hasRef()
     // reports false (the disentangle analogue of the close() reset above).
@@ -497,6 +483,7 @@ void MessagePort::onDidChangeListenerImpl(EventTarget& self, const AtomString& e
         return;
 
     auto& port = static_cast<MessagePort&>(self);
+    bool hadListeners = port.m_messageEventCount > 0;
     switch (kind) {
     case Add:
         port.m_messageEventCount++;
@@ -510,6 +497,10 @@ void MessagePort::onDidChangeListenerImpl(EventTarget& self, const AtomString& e
         break;
     }
     port.updateListenerEventLoopRef();
+    // node's setupPortReferencing unref()s the port outright when its last 'message'
+    // listener goes away, so an earlier .ref() (or onmessage=) does not outlive it.
+    if (hadListeners && port.m_messageEventCount == 0)
+        port.releaseJsRef();
 }
 
 bool MessagePort::addEventListener(const AtomString& eventType, Ref<EventListener>&& listener, const AddEventListenerOptions& options)
@@ -576,19 +567,28 @@ void MessagePort::jsRef(JSGlobalObject* lexicalGlobalObject)
     }
 }
 
-void MessagePort::jsUnref(JSGlobalObject* lexicalGlobalObject)
+void MessagePort::jsUnref()
 {
     // Also release the listener loop-ref; otherwise an always-listening transferred
     // port (a postMessageToThread control port) would pin the event loop forever.
     if (m_isRefd) {
         m_isRefd = false;
         updateListenerEventLoopRef();
     }
-    if (m_hasRef) {
-        m_hasRef = false;
-        deref();
-        Bun__eventLoop__refKeepAlive(WebCore::clientData(lexicalGlobalObject->vm())->bunVM, -1);
-    }
+    releaseJsRef();
+}
+
+void MessagePort::releaseJsRef()
+{
+    if (!m_hasRef)
+        return;
+    m_hasRef = false;
+    // Same per-thread VM that jsRef() ref'd through the lexical global. The context is
+    // still attached here: contextDestroyed() reaches this via close() before detaching.
+    if (auto* context = scriptExecutionContext())
+        context->unrefEventLoop();
+    // Drops the self-ref, so `this` must not be touched afterwards.
+    deref();
 }
 
 } // namespace WebCore

src/jsc/bindings/webcore/MessagePort.h

@@ -114,7 +114,7 @@ class MessagePort final : public ActiveDOMObject, public EventTarget, public Thr
     JSValue tryTakeMessage(JSGlobalObject*, bool& hadMessage);
 
     void jsRef(JSGlobalObject*);
-    void jsUnref(JSGlobalObject*);
+    void jsUnref();
     // Report the actual loop-ref state (matches Node's uv_has_ref), not the intent flag.
     bool jsHasRef() { return m_hasRef || m_listenerLoopRefActive; }
 
@@ -159,7 +159,11 @@ class MessagePort final : public ActiveDOMObject, public EventTarget, public Thr
     // Read from the GC thread: a port whose only listener is 'close' must survive
     // until that event is delivered, or the peer's close is lost to a collection.
     std::atomic<bool> m_hasCloseEventListener { false };
+    // jsRef() (.ref(), or a callable .onmessage=) took a self-ref plus an event-loop ref.
+    // Released by unref(), close(), the peer closing, a transfer, and (as node's
+    // setupPortReferencing does) the removal of the last 'message' listener.
     bool m_hasRef { false };
+    void releaseJsRef();
 
     // Whether .ref()/.unref() want this port to keep the loop alive (default refd);
     // independent of m_hasRef (the .onmessage=/.ref() keepalive).

test/js/node/worker_threads/worker_threads.test.ts

@@ -997,6 +997,202 @@ test("hasRef() survives collection of the unreferenced peer", () => {
   port1.close();
 });
 
+// port1.hasRef() after running `scenario` on a fresh channel. port2 stays reachable
+// until the end so that collecting it cannot close port1 underneath the scenario.
+async function hasRefAfter(scenario: (port1: MessagePort, port2: MessagePort) => void | Promise<void>) {
+  const { port1, port2 } = new MessageChannel();
+  try {
+    await scenario(port1, port2);
+    return port1.hasRef();
+  } finally {
+    port1.close();
+    port2.close();
+  }
+}
+
+// node's setupPortReferencing (lib/internal/worker/io.js) unref()s a port when its
+// last 'message' listener is removed, whether or not ref() was called on it before.
+test("removing the last 'message' listener releases an explicit ref() as well", async () => {
+  const f = () => {};
+  const g = () => {};
+  const results = {
+    refThenListenerRemoved: await hasRefAfter(p => {
+      p.ref();
+      p.on("message", f);
+      p.off("message", f);
+    }),
+    listenerThenRefThenRemoved: await hasRefAfter(p => {
+      p.on("message", f);
+      p.ref();
+      p.off("message", f);
+    }),
+    viaRemoveEventListener: await hasRefAfter(p => {
+      p.ref();
+      p.addEventListener("message", f);
+      p.removeEventListener("message", f);
+    }),
+    viaOnceListenerFiring: await hasRefAfter(async (p, peer) => {
+      p.ref();
+      const fired = new Promise<void>(resolve => p.once("message", () => resolve()));
+      peer.postMessage(1);
+      await fired;
+    }),
+    // None of these remove the last 'message' listener, so nothing is released.
+    oneOfTwoRemoved: await hasRefAfter(p => {
+      p.ref();
+      p.on("message", f);
+      p.on("message", g);
+      p.off("message", f);
+    }),
+    neverAddedListenerRemoved: await hasRefAfter(p => {
+      p.ref();
+      p.off("message", f);
+    }),
+    closeListenerRemoved: await hasRefAfter(p => {
+      p.ref();
+      p.on("close", f);
+      p.off("close", f);
+    }),
+    // The release is not sticky: a later listener or ref() refs the port again.
+    listenerAddedAgain: await hasRefAfter(p => {
+      p.ref();
+      p.on("message", f);
+      p.off("message", f);
+      p.on("message", g);
+    }),
+    refCalledAgain: await hasRefAfter(p => {
+      p.ref();
+      p.on("message", f);
+      p.off("message", f);
+      p.ref();
+    }),
+  };
+  expect(results).toEqual({
+    refThenListenerRemoved: false,
+    listenerThenRefThenRemoved: false,
+    viaRemoveEventListener: false,
+    viaOnceListenerFiring: false,
+    oneOfTwoRemoved: true,
+    neverAddedListenerRemoved: true,
+    closeListenerRemoved: true,
+    listenerAddedAgain: true,
+    refCalledAgain: true,
+  });
+});
+
+// bun's removeAllListeners() removes the listeners one by one, so it releases the port
+// exactly like off() on each of them would. (node's NodeEventTarget#removeAllListeners
+// bypasses its listener hooks and leaves a port with no listeners ref'd; bun does not
+// keep the listener-count ref there either, and the explicit ref follows the same rule.)
+test("removeAllListeners() releases an explicit ref() like off() does", async () => {
+  const f = () => {};
+  const results = {
+    byType: await hasRefAfter(p => {
+      p.ref();
+      p.on("message", f);
+      p.removeAllListeners("message");
+    }),
+    all: await hasRefAfter(p => {
+      p.ref();
+      p.on("message", f);
+      p.removeAllListeners();
+    }),
+    unrelatedType: await hasRefAfter(p => {
+      p.ref();
+      p.on("message", f);
+      p.removeAllListeners("close");
+    }),
+  };
+  expect(results).toEqual({ byType: false, all: false, unrelatedType: true });
+});
+
+// node: assigning a non-function to onmessage removes the handler, and the port is
+// unref()'d only if that removed its last 'message' listener. Clearing a handler that
+// was never set, or while on() listeners remain, leaves the ref state alone.
+test("onmessage = null releases the port only when it removed the last 'message' listener", async () => {
+  const f = () => {};
+  const g = () => {};
+  const results = {
+    handlerCleared: await hasRefAfter(p => {
+      p.onmessage = f;
+      p.onmessage = null;
+    }),
+    refdHandlerCleared: await hasRefAfter(p => {
+      p.ref();
+      p.onmessage = f;
+      p.onmessage = null;
+    }),
+    nothingToClear: await hasRefAfter(p => {
+      p.ref();
+      p.onmessage = null;
+    }),
+    onListenerRemains: await hasRefAfter(p => {
+      p.on("message", f);
+      p.onmessage = g;
+      p.onmessage = null;
+    }),
+    onListenerOnly: await hasRefAfter(p => {
+      p.on("message", f);
+      p.onmessage = null;
+    }),
+  };
+  expect(results).toEqual({
+    handlerCleared: false,
+    refdHandlerCleared: false,
+    nothingToClear: true,
+    onListenerRemains: true,
+    onListenerOnly: true,
+  });
+});
+
+// End to end: once the listener is gone nothing about the port may hold the process open.
+test.concurrent("a ref()'d port whose last 'message' listener was removed lets the process exit", async () => {
+  await using proc = Bun.spawn({
+    cmd: [
+      bunExe(),
+      "-e",
+      `const { MessageChannel } = require("worker_threads");
+       const { port1, port2 } = new MessageChannel();
+       // Keep the peer reachable: collecting it would close port1 (and release its refs) by itself.
+       globalThis.keep = port2;
+       const f = () => {};
+       port1.ref();
+       port1.on("message", f);
+       port1.off("message", f);
+       console.log("hasRef=" + port1.hasRef());
+       // Never fires when the port released the loop; bounds the failure mode (a process
+       // that stays alive) instead of letting the test run into its timeout.
+       setTimeout(() => { console.log("still alive"); process.exit(1); }, 2_000).unref();`,
+    ],
+    env: bunEnv,
+    stdout: "pipe",
+    stderr: "inherit",
+  });
+  const [stdout, exitCode] = await Promise.all([proc.stdout.text(), proc.exited]);
+  expect(stdout).toBe("hasRef=false\n");
+  expect(exitCode).toBe(0);
+});
+
+// Same rule for parentPort: node's worker exits once the listeners are gone even if
+// the script ref()'d the port explicitly.
+test("a ref()'d parentPort whose last 'message' listener was removed lets the worker exit", async () => {
+  const w = new Worker(
+    `const { parentPort } = require("worker_threads");
+     const f = () => {};
+     parentPort.ref();
+     parentPort.on("message", f);
+     parentPort.off("message", f);
+     parentPort.postMessage(parentPort.hasRef());
+     // Never fires when the port released the loop; bounds the failure mode (a worker
+     // that stays alive) with a distinctive exit code.
+     setTimeout(() => process.exit(7), 2_000).unref();`,
+    { eval: true },
+  );
+  const hasRef = once(w, "message");
+  const exited = once(w, "exit");
+  expect({ hasRef: (await hasRef)[0], exitCode: (await exited)[0] }).toEqual({ hasRef: false, exitCode: 0 });
+});
+
 // markAsUncloneable blocks *cloning*, not transfer: a marked port in the transfer
 // list is moved, so node lets it through and it still works on the far side.
 test("markAsUncloneable blocks cloning a port but not transferring it", async () => {