Fire the close event on MessagePort

src/jsc/bindings/webcore/MessagePort.cpp

@@ -28,6 +28,7 @@
 #include "MessagePort.h"
 
 #include "BunClientData.h"
+#include "Event.h"
 #include "EventNames.h"
 #include "MessageEvent.h"
 #include "MessagePortPipe.h"
@@ -110,9 +111,15 @@
 
     // m_pipe is held for the port's whole lifetime (the GC thread reads
     // it in hasPendingActivity()); marking our side Closed is sufficient.
+    // close() also wakes the entangled peer so it can fire its own 'close'
+    // event after draining any queued messages.
     m_pipe->close(m_side);
 
-    removeAllEventListeners();
+    // Fire our own 'close' event asynchronously (Node + HTML semantics), then
+    // tear down listeners. If there is no close listener or JS can't run
+    // (context teardown), tear down synchronously as before. scheduleCloseEvent()
+    // takes its own strong ref, so it is safe to run before releasing m_hasRef.
+    bool scheduledClose = scheduleCloseEvent();
 
     // Release the self-reference taken by jsRef() (set when .onmessage is
     // assigned or .ref() is called from JS). The JS .close() binding calls
@@ -127,6 +134,103 @@
             context->unrefEventLoop();
         deref();
     }
+
+    if (!scheduledClose) {
+        // No 'close' event will be dispatched for this port (no listener, or
+        // JS can't run during teardown). Mark the close consumed so that a
+        // close listener added after close() does not pin the already-closed
+        // wrapper forever via hasPendingActivity().
+        m_closeEventDispatched = true;
+        removeAllEventListeners();
+    }
+}
+
+void MessagePort::startForClose()
+{
+    // Register with the pipe so MessagePortPipe::close() on the peer can
+    // schedule a drain that reaches us and dispatches 'close'. Unlike start(),
+    // this does not set m_started, so a later 'message' listener still runs
+    // start() and re-attaches to flush any buffered messages. attach() is
+    // idempotent, so calling it again from start() is harmless.
+    if (!isEntangled())
+        return;
+    auto* context = scriptExecutionContext();
+    if (!context)
+        return;
+    m_pipe->attach(m_side, context->identifier(), ThreadSafeWeakPtr<MessagePort> { *this });
+}
+
+void MessagePort::dispatchCloseEvent()
+{
+    if (m_closeEventDispatched)
+        return;
+    m_closeEventDispatched = true;
+
+    auto* context = scriptExecutionContext();
+    if (!context || !context->globalObject())
+        return;
+
+    auto* globalObject = defaultGlobalObject(context->globalObject());
+    if (Zig::GlobalObject::scriptExecutionStatus(globalObject, globalObject) != ScriptExecutionStatus::Running)
+        return;
+
+    // Bypass MessagePort::dispatchEvent()'s detached guard: by the time the
+    // close task runs the port is already detached, but the 'close' event must
+    // still reach its listener.
+    EventTarget::dispatchEvent(Event::create(eventNames().closeEvent, Event::CanBubble::No, Event::IsCancelable::No));
+}
+
+void MessagePort::dispatchCloseEventSelf()
+{
+    dispatchCloseEvent();
+    removeAllEventListeners();
+    m_hasMessageEventListener = false;
+    m_hasCloseEventListener = false;
+}
+
+bool MessagePort::scheduleCloseEvent()
+{
+    if (m_closeEventDispatched || !m_hasCloseEventListener)
+        return false;
+
+    auto* context = scriptExecutionContext();
+    if (!context || !context->globalObject())
+        return false;
+
+    auto* globalObject = defaultGlobalObject(context->globalObject());
+    if (Zig::GlobalObject::scriptExecutionStatus(globalObject, globalObject) != ScriptExecutionStatus::Running)
+        return false;
+
+    return ScriptExecutionContext::postTaskTo(context->identifier(), [protectedThis = Ref { *this }](ScriptExecutionContext&) {
+        protectedThis->dispatchCloseEventSelf();
+    });
+}
+
+void MessagePort::dispatchCloseEventFromPeer()
+{
+    if (m_isDetached || m_closeEventDispatched || !m_hasCloseEventListener)
+        return;
+
+    // Runs JS (the close handler), which may drop the last external ref.
+    Ref protectedThis { *this };
+    // Stop message delivery and make a re-entrant close() a no-op. The close
+    // branch of hasPendingActivity() (driven by pipe state + the listener
+    // flag) keeps the wrapper alive across the dispatch below.
+    m_isDetached = true;
+
+    dispatchCloseEvent();
+
+    m_pipe->close(m_side);
+    removeAllEventListeners();
+    m_hasMessageEventListener = false;
+    m_hasCloseEventListener = false;
+
+    if (m_hasRef) {
+        m_hasRef = false;
+        if (auto* context = scriptExecutionContext())
+            context->unrefEventLoop();
+        deref();
+    }
 }
 
 TransferredMessagePort MessagePort::disentangle()
@@ -244,12 +348,31 @@
     // atomic loads. The plain bool reads can observe stale values but
     // cannot crash — at worst the wrapper is collected one cycle early
     // or late, which is the same tolerance as before this refactor.
-    if (!scriptExecutionContext() || m_isDetached)
+    if (!scriptExecutionContext())
+        return false;
+
+    uint64_t s = m_pipe->state(m_side);
+
+    // Keep the wrapper (and its 'close' listener) alive until a pending close
+    // event is dispatched. A close is pending, and will actually fire, when a
+    // close listener is registered, it has not been dispatched yet, and either:
+    //   - this side has closed (the closing port's own scheduled 'close'), or
+    //   - this side is still attached and the peer has closed (the drain will
+    //     dispatch 'close' to us).
+    // The !m_closeEventDispatched guard is what lets an already-closed port be
+    // collected: close() marks it dispatched when no listener will fire, so a
+    // close listener added afterwards cannot pin the wrapper forever. The
+    // Attached guard avoids pinning a never-started port whose peer closed.
+    if (m_hasCloseEventListener && !m_closeEventDispatched
+        && ((s & MessagePortPipe::Closed)
+            || ((s & MessagePortPipe::Attached) && !m_pipe->isOtherSideOpen(m_side))))
+        return true;
+
+    if (m_isDetached)
         return false;
     if (!m_hasMessageEventListener)
         return false;
 
-    uint64_t s = m_pipe->state(m_side);
     // Keep alive if there are messages already queued for us, or the peer
     // is still open and could send more.
     return MessagePortPipe::queuedCount(s) > 0 || m_pipe->isOtherSideOpen(m_side);
@@ -312,6 +435,9 @@
     if (eventType == eventNames().messageEvent) {
         start();
         m_hasMessageEventListener = true;
+    } else if (eventType == eventNames().closeEvent) {
+        startForClose();
+        m_hasCloseEventListener = true;
     }
     return EventTarget::addEventListener(eventType, WTF::move(listener), options);
 }
@@ -321,6 +447,8 @@
     auto result = EventTarget::removeEventListener(eventType, listener, options);
     if (!hasEventListeners(eventNames().messageEvent))
         m_hasMessageEventListener = false;
+    if (!hasEventListeners(eventNames().closeEvent))
+        m_hasCloseEventListener = false;
     return result;
 }
 

src/jsc/bindings/webcore/MessagePort.h

@@ -84,6 +84,16 @@ class MessagePort final : public ContextDestructionObserver, public EventTarget,
     // Called by the pipe on this port's context thread with one dequeued message.
     void dispatchOneMessage(ScriptExecutionContext&, MessageWithMessagePorts&&);
 
+    // Called by the pipe drain on this port's context thread once the inbox is
+    // empty and the entangled peer has closed: fires 'close' (if a listener is
+    // registered) and tears the port down.
+    void dispatchCloseEventFromPeer();
+
+    // The pipe drain delivers queued messages only to a port that is listening
+    // for them; otherwise messages stay buffered (a port attached solely for a
+    // 'close' listener must not drop them). Matches Node/HTML start() semantics.
+    bool isListeningForMessages() const { return m_hasMessageEventListener; }
+
     // Only here for JSMessagePortCustom's GC optimization; always null.
     MessagePort* locallyEntangledPort() { return nullptr; }
 
@@ -117,6 +127,20 @@ class MessagePort final : public ContextDestructionObserver, public EventTarget,
 
     void contextDestroyed() final;
 
+    // Attach to the pipe so the peer can wake this port to dispatch 'close',
+    // without enabling message delivery (unlike start(), which a 'message'
+    // listener calls). Idempotent with a later start().
+    void startForClose();
+
+    // Fires the 'close' event on this port (once). Safe to call after the port
+    // is detached: it bypasses the detached guard in dispatchEvent().
+    void dispatchCloseEvent();
+    // Task body for the closing port's own asynchronous 'close' event.
+    void dispatchCloseEventSelf();
+    // Schedules dispatchCloseEventSelf() on this port's context. Returns true
+    // if a task was posted (in which case listener teardown is deferred to it).
+    bool scheduleCloseEvent();
+
     bool isEntangled() const { return !m_isDetached; }
 
     // Held for the port's entire lifetime — never nulled — so that the GC
@@ -128,6 +152,8 @@ class MessagePort final : public ContextDestructionObserver, public EventTarget,
     bool m_started { false };
     bool m_isDetached { false };
     bool m_hasMessageEventListener { false };
+    bool m_hasCloseEventListener { false };
+    bool m_closeEventDispatched { false };
     bool m_hasRef { false };
 
     uint32_t m_messageEventCount { 0 };

src/jsc/bindings/webcore/MessagePortPipe.cpp

@@ -98,7 +98,10 @@
     auto& s = m_sides[side];
 
     RefPtr<MessagePort> port;
-    size_t limit;
+    size_t limit = 0;
+    // Set when the inbox empties while the port is still attached: the port may
+    // then owe a 'close' event (dispatched below, outside the lock).
+    bool emptied = false;
     {
         Locker locker { s.lock };
         // This task was posted to `expectedCtx` (and is running there). If
@@ -112,9 +115,18 @@
         uint64_t st = s.state.load(std::memory_order_relaxed);
         if (!port || s.inbox.isEmpty()) {
             s.state.store(st & ~DrainScheduled, std::memory_order_release);
-            return;
+            emptied = port && (st & Attached) != 0;
+        } else {
+            limit = std::max<size_t>(s.inbox.size(), 1000);
         }
-        limit = std::max<size_t>(s.inbox.size(), 1000);
+    }
+
+    if (!limit) {
+        // Nothing queued. If our peer has closed, the attached port still owes
+        // a 'close' event.
+        if (emptied && !isOtherSideOpen(side))
+            port->dispatchCloseEventFromPeer();
+        return;
     }
 
     auto* context = port->scriptExecutionContext();
@@ -141,9 +153,20 @@
                 break;
             uint64_t st = s.state.load(std::memory_order_relaxed);
             if (!(st & Attached) || s.inbox.isEmpty()) {
+                s.state.store(st & ~DrainScheduled, std::memory_order_release);
+                emptied = (st & Attached) != 0;
+                break;
+            }
+            // A port attached only for a 'close' listener (no 'message'
+            // listener) must not have its queued messages dispatched to no one.
+            // Leave them buffered; adding a 'message' listener runs start(),
+            // which re-attaches and reschedules this drain. Don't set `emptied`
+            // — the inbox is non-empty, so a pending 'close' must wait behind
+            // the undelivered messages (matching Node).
+            if (!port->isListeningForMessages()) {
                 s.state.store(st & ~DrainScheduled, std::memory_order_release);
                 break;
             }
             if (limit-- == 0) {
                 // Yield to the rest of the event loop; DrainScheduled stays
                 // set so concurrent sends don't double-schedule.
@@ -165,6 +188,10 @@
 
     if (rescheduleCtx)
         scheduleDrain(side, rescheduleCtx);
+    else if (emptied && !isOtherSideOpen(side))
+        // Inbox fully drained and the peer has closed: deliver 'close' after
+        // all queued messages, matching Node's ordering.
+        port->dispatchCloseEventFromPeer();
 }
 
 std::optional<MessageWithMessagePorts> MessagePortPipe::takeOne(uint8_t side)
@@ -189,7 +216,11 @@
         s.port = WTF::move(port);
         uint64_t st = s.state.load(std::memory_order_relaxed);
         uint64_t ns = (st | Attached) & ~Closed;
-        if (queuedCount(st) > 0 && !(st & DrainScheduled)) {
+        // Schedule a drain if there is work to dispatch: queued messages, or a
+        // peer that has already closed (a 'close' event is owed). The latter
+        // covers a 'close' listener added after the peer closed, which would
+        // otherwise have missed the peer's wake-up in close().
+        if (!(st & DrainScheduled) && (queuedCount(st) > 0 || !isOtherSideOpen(side))) {
             ns |= DrainScheduled;
             wakeCtx = ctxId;
         }
@@ -241,6 +272,27 @@
             dropped = std::exchange(s.inbox, {});
         }
 
+        // Wake the entangled peer (after its queued messages drain) so it can
+        // dispatch a 'close' event now that this side has closed. Only needed
+        // when the peer is attached and no drain is already in flight; an
+        // in-flight drain observes the Closed bit we just stored (set before
+        // this check) and dispatches close itself. Marked Closed above →
+        // checked here preserves that ordering.
+        {
+            auto& peer = pipe->m_sides[1 - sd];
+            ScriptExecutionContextIdentifier peerCtx = 0;
+            {
+                Locker locker { peer.lock };
+                uint64_t ps = peer.state.load(std::memory_order_relaxed);
+                if ((ps & Attached) && !(ps & DrainScheduled)) {
+                    peer.state.store(ps | DrainScheduled, std::memory_order_release);
+                    peerCtx = peer.ctxId;
+                }
+            }
+            if (peerCtx)
+                pipe->scheduleDrain(1 - sd, peerCtx);
+        }
+
         // Harvest transferred pipes before `dropped` destructs so their
         // ~TransferredMessagePort sees pipe == nullptr and is a no-op.
         for (auto& message : dropped) {