MessagePort: fire 'close' on a port that is transferred away

src/jsc/bindings/webcore/MessagePort.cpp

@@ -241,9 +241,12 @@ void MessagePort::close()
         updateListenerEventLoopRef();
     }
 
-    // Defer 'close' to a task (node fires it at uv close-callback timing, i.e.
-    // after sync code and microtasks), so a listener added after close() still
-    // observes it and close(cb) interleaves with other listeners.
+    queueCloseEvent();
+}
+
+void MessagePort::queueCloseEvent()
+{
+    // A task rather than synchronous: node fires it from the uv close callback, so a listener added right after close() or the transfer still sees it.
     if (isContextStopped()) {
         removeAllEventListeners();
         return;
@@ -295,17 +298,10 @@ TransferredMessagePort MessagePort::disentangle()
 {
     ASSERT(isEntangled());
 
-    // Drop any message listeners (and the event-loop ref they carry) while
-    // this port is still attached to its context; after observeContext(null)
-    // there would be nothing to unref.
-    removeAllEventListeners();
+    // The listeners themselves are dropped by the close task queued below.
     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
-    // 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.
+    // Release jsRef()'s self-reference here, since close() no-ops on a detached port; the caller's RefPtr keeps us alive across deref().
     if (m_hasRef) {
         m_hasRef = false;
         if (auto* context = scriptExecutionContext())
@@ -329,12 +325,8 @@ TransferredMessagePort MessagePort::disentangle()
     m_isDetached = true;
     m_started = false;
 
-    // We can't receive any messages or generate any events after this, so remove ourselves from the list of active ports.
-    if (auto* context = scriptExecutionContext()) {
-        context->willDestroyActiveDOMObject(*this);
-        context->willDestroyDestructionObserver(*this);
-    }
-    observeContext(nullptr);
+    // Stays attached to the context like a close()d port: the task dispatches through it, and only pins the wrapper while we have one.
+    queueCloseEvent();
 
     return TransferredMessagePort { m_pipe.copyRef(), m_side };
 }
@@ -373,6 +365,7 @@ void MessagePort::dispatchOneMessage(ScriptExecutionContext& context, MessageWit
 JSValue MessagePort::tryTakeMessage(JSGlobalObject* lexicalGlobalObject, bool& hadMessage)
 {
     hadMessage = false;
+    // Also what stops a transferred-away object, which keeps its context, from taking from the pipe side its receiver now owns.
     if (!isEntangled())
         return jsUndefined();
 

src/jsc/bindings/webcore/MessagePort.h

@@ -84,7 +84,7 @@ class MessagePort final : public ActiveDOMObject, public EventTarget, public Thr
     void peerClosed();
     void dispatchCloseEvent();
 
-    // Transfer machinery.
+    // Transfer machinery. disentangle() also queues 'close' on the transferred-away object, like node's TransferForMessaging(): https://github.com/nodejs/node/blob/v26.3.0/src/node_messaging.cc#L921-L924
     static ExceptionOr<Vector<TransferredMessagePort>> disentanglePorts(Vector<RefPtr<MessagePort>>&&);
     static Vector<RefPtr<MessagePort>> entanglePorts(ScriptExecutionContext&, Vector<TransferredMessagePort>&&);
     static Ref<MessagePort> entangle(ScriptExecutionContext&, TransferredMessagePort&&);
@@ -131,6 +131,8 @@ class MessagePort final : public ActiveDOMObject, public EventTarget, public Thr
 
     // Deliver messages already queued when close() is called, before teardown.
     void flushQueuedMessagesBeforeClose();
+    // Shared tail of close() and disentangle(): fires 'close' from a task, then drops the listeners.
+    void queueCloseEvent();
 
     bool isEntangled() const { return !m_isDetached; }
 

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

@@ -1223,6 +1223,138 @@ test("dropping a transferred port notifies its peer", async () => {
   port1.close();
 });
 
+// Transferring a port closes the sender's object (node closes its handle), so that object
+// fires 'close' once, asynchronously. The channel itself moves to the receiver: the peer
+// sees nothing, and the channel's eventual close goes to the received port.
+describe("a port transferred away fires 'close' on the sender's object", () => {
+  test("postMessage to a local port", async () => {
+    const { port1, port2 } = new MessageChannel();
+    const carrier = new MessageChannel();
+    const events: string[] = [];
+    port1.on("close", () => events.push("port1 close (listener added before the transfer)"));
+    port2.on("close", () => events.push("port2 close"));
+    carrier.port1.postMessage(port1, [port1]);
+    port1.on("close", () => events.push("port1 close (listener added after the transfer)"));
+    events.push("postMessage returned");
+    // The event is queued as a task, which runs before the first setImmediate. A fixed
+    // window instead of awaiting it keeps the exact lists below meaningful: they also
+    // assert what must not fire (port2 now, port1 again at the end).
+    for (let i = 0; i < 2; i++) await new Promise(r => setImmediate(r));
+    expect(events).toEqual([
+      "postMessage returned",
+      "port1 close (listener added before the transfer)",
+      "port1 close (listener added after the transfer)",
+    ]);
+
+    // The channel is intact and belongs to the received port: the transferred-away object
+    // shares the pipe side with it (and, like a closed port, still has a context), so it
+    // must not be able to take the receiver's messages. (Node returns undefined here too,
+    // once the port's 'close' has fired.)
+    port2.postMessage("to the received port");
+    expect(receiveMessageOnPort(port1)).toBeUndefined();
+    const received = await new Promise<MessagePort>(resolve => carrier.port2.once("message", resolve));
+    expect(await new Promise(resolve => received.once("message", resolve))).toBe("to the received port");
+    received.postMessage("from the received port");
+    expect(await new Promise(resolve => port2.once("message", resolve))).toBe("from the received port");
+
+    // Closing the channel now closes port2 and the received port; the object that was
+    // transferred away does not fire a second time.
+    received.on("message", () => {});
+    const receivedClosed = new Promise<void>(resolve => received.once("close", () => resolve()));
+    port2.close();
+    await receivedClosed;
+    for (let i = 0; i < 2; i++) await new Promise(r => setImmediate(r));
+    expect(events).toEqual([
+      "postMessage returned",
+      "port1 close (listener added before the transfer)",
+      "port1 close (listener added after the transfer)",
+      "port2 close",
+    ]);
+    carrier.port1.close();
+    carrier.port2.close();
+  });
+
+  // The transfer is the last thing the script does: the event still fires, and the
+  // pending event does not keep the process alive afterwards.
+  test("the event fires before the process exits and does not keep it alive", async () => {
+    await using proc = Bun.spawn({
+      cmd: [
+        bunExe(),
+        "-e",
+        `const { MessageChannel } = require("worker_threads");
+         const { port1 } = new MessageChannel();
+         const carrier = new MessageChannel();
+         port1.on("close", () => console.log("port1 close"));
+         carrier.port1.postMessage(port1, [port1]);`,
+      ],
+      env: bunEnv,
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    expect({ stdout: stdout.trim(), stderr, exitCode, signalCode: proc.signalCode }).toEqual({
+      stdout: "port1 close",
+      stderr: "",
+      exitCode: 0,
+      signalCode: null,
+    });
+  });
+
+  // Ports transferred through the Worker constructor's transferList, through
+  // worker.postMessage(), and from the worker through parentPort.postMessage(). Each
+  // route also passes a message through the port the other side received, so the
+  // sender's 'close' is shown to leave the channel itself intact. The process exits on
+  // its own once the worker has finished, so every line is in before it does.
+  test("transfers to and from a Worker", async () => {
+    await using proc = Bun.spawn({
+      cmd: [
+        bunExe(),
+        "-e",
+        `const { Worker, MessageChannel } = require("worker_threads");
+         const viaConstructor = new MessageChannel();
+         const viaPostMessage = new MessageChannel();
+         viaConstructor.port2.on("close", () => console.log("transferList: close"));
+         viaPostMessage.port2.on("close", () => console.log("worker.postMessage: close"));
+         viaConstructor.port1.once("message", m => console.log("transferList: " + m));
+         viaPostMessage.port1.once("message", m => console.log("worker.postMessage: " + m));
+         const w = new Worker(
+           \`const { parentPort, workerData, MessageChannel } = require("worker_threads");
+            workerData.postMessage("received port works");
+            parentPort.once("message", received => {
+              received.postMessage("received port works");
+              const { port1, port2 } = new MessageChannel();
+              port2.postMessage("received port works");
+              port1.on("close", () => parentPort.postMessage("parentPort.postMessage: close"));
+              parentPort.postMessage(port1, [port1]);
+            });\`,
+           { eval: true, workerData: viaConstructor.port2, transferList: [viaConstructor.port2] },
+         );
+         w.postMessage(viaPostMessage.port2, [viaPostMessage.port2]);
+         w.on("message", m => {
+           if (typeof m === "string") console.log(m);
+           else m.once("message", x => console.log("parentPort.postMessage: " + x));
+         });`,
+      ],
+      env: bunEnv,
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    // The routes complete independently of each other, so the order of the lines varies.
+    expect({ lines: stdout.trim().split("\n").sort(), stderr, exitCode, signalCode: proc.signalCode }).toEqual({
+      lines: [
+        "parentPort.postMessage: close",
+        "parentPort.postMessage: received port works",
+        "transferList: close",
+        "transferList: received port works",
+        "worker.postMessage: close",
+        "worker.postMessage: received port works",
+      ],
+      stderr: "",
+      exitCode: 0,
+      signalCode: null,
+    });
+  });
+});
+
 // close() outside a dispatch drops whatever is queued; close() from inside a
 // 'message' handler lets the in-flight drain finish. Both are node's behaviour.
 test("close() drops queued messages unless it runs inside a dispatch", async () => {

test/js/web/workers/message-channel.test.ts

@@ -1,3 +1,5 @@
+import { heapStats } from "bun:jsc";
+
 test("simple usage", done => {
   const channel = new MessageChannel();
   const port1 = channel.port1;
@@ -352,6 +354,38 @@ test("a pending close event survives GC after the port becomes unreachable", asy
   expect(fired).toBe(50);
 });
 
+// A transfer closes the sender's port object the same way (node fires 'close' on it), so
+// disentangle() queues the same task, and that object must stay attached to its context
+// until the task runs: the pending activity only pins the wrapper while it has a context.
+test("a transferred-away port's pending close event survives GC after the port becomes unreachable", async () => {
+  const count = () => heapStats().objectTypeCounts.MessagePort ?? 0;
+  Bun.gc(true);
+  const base = count();
+  let fired = 0;
+  const carrier = new MessageChannel();
+  for (let i = 0; i < 50; i++) {
+    (() => {
+      const { port1 } = new MessageChannel();
+      port1.addEventListener("close", () => fired++);
+      carrier.port1.postMessage(port1, [port1]);
+    })();
+    if (i % 10 === 0) Bun.gc(true);
+  }
+  Bun.gc(true);
+  // The events are queued as tasks, which run before the first setImmediate; the fixed
+  // window (rather than awaiting the 50th event) is what lets the exact count below also
+  // catch a port firing twice.
+  for (let i = 0; i < 4; i++) await new Promise(r => setImmediate(r));
+  expect(fired).toBe(50);
+  // Once the event has fired nothing pins the transferred-away objects (or their
+  // unreferenced peers) any more; only the carrier pair is still reachable.
+  Bun.gc(true);
+  Bun.gc(true);
+  expect(count() - base).toBeLessThanOrEqual(2 + 10);
+  carrier.port1.close();
+  carrier.port2.close();
+});
+
 // The peer's notifyPeerClosed() task only holds a weak ref back to this port, so a
 // port whose only listener is 'close' must survive GC until the event is delivered.
 test("a close event from the peer survives GC of the unreachable port", async () => {
@@ -375,7 +409,6 @@ test("a close event from the peer survives GC of the unreachable port", async ()
 // which never collects an entangled port). Bound the retention: explicitly-closed pairs
 // must still be swept, so a regression that leaks closed ports too would show as growth.
 test("explicitly-closed close-listener ports are collected; open ones are pinned like Node", async () => {
-  const { heapStats } = require("bun:jsc");
   const count = () => heapStats().objectTypeCounts.MessagePort ?? 0;
   for (let i = 0; i < 4; i++) await new Promise(r => setImmediate(r));
   Bun.gc(true);

test/js/web/workers/worker-postmessage-transfer.test.ts

@@ -136,4 +136,61 @@ describe("self.postMessage transfer list", () => {
       URL.revokeObjectURL(url);
     }
   });
+
+  // Transferring a port closes the sender's object (node fires 'close' on it, asynchronously)
+  // and nothing else: the peer stays open and the receiver's port talks to it. Worker#postMessage
+  // and the worker-side self.postMessage each disentangle the port themselves, so both
+  // directions are checked. The event is a task, which runs before the first setImmediate;
+  // both sides record for two of them rather than awaiting the event so that the recorded
+  // lists also show anything that fired when it should not have.
+  test("a port transferred with postMessage fires 'close' on the sender's object, in both directions", async () => {
+    const url = URL.createObjectURL(
+      new Blob([
+        `
+          const { port1, port2 } = new MessageChannel();
+          port2.postMessage("queued before the transfer"); // travels with port1
+          const events = [];
+          port1.addEventListener("close", () => events.push("close"));
+          port2.addEventListener("close", () => events.push("peer close"));
+          self.postMessage(port1, [port1]);
+          events.push("postMessage returned");
+          setImmediate(() => setImmediate(() => self.postMessage(events)));
+          self.onmessage = e => e.data.postMessage("sent through the received port");
+        `,
+      ]),
+    );
+    const worker = new Worker(url);
+    try {
+      const fromWorker: any[] = [];
+      const workerEvents = new Promise<string[]>((resolve, reject) => {
+        worker.onerror = e => reject(e.error ?? e.message ?? e);
+        worker.onmessage = e => {
+          fromWorker.push(e.data);
+          if (Array.isArray(e.data)) resolve(e.data);
+        };
+      });
+
+      const { port1, port2 } = new MessageChannel();
+      const events: string[] = [];
+      port1.addEventListener("close", () => events.push("close"));
+      port2.addEventListener("close", () => events.push("peer close"));
+      const throughPort2 = new Promise<string>(resolve => (port2.onmessage = e => resolve(e.data)));
+      worker.postMessage(port1, [port1]);
+      events.push("postMessage returned");
+      await new Promise(r => setImmediate(() => setImmediate(r)));
+      expect(events).toEqual(["postMessage returned", "close"]);
+      expect(await throughPort2).toBe("sent through the received port");
+
+      expect(await workerEvents).toEqual(["postMessage returned", "close"]);
+      const received = fromWorker[0];
+      expect(received).toBeInstanceOf(MessagePort);
+      const throughReceived = new Promise<string>(resolve => (received.onmessage = e => resolve(e.data)));
+      expect(await throughReceived).toBe("queued before the transfer");
+      received.close();
+      port2.close();
+    } finally {
+      worker.terminate();
+      URL.revokeObjectURL(url);
+    }
+  });
 });