Stop dispatching parent messages to self.onmessage in node:worker_threads workers

src/js/node/worker_threads.ts

@@ -129,22 +129,270 @@
 
 // TODO: parent port emulation is not complete
 function fakeParentPort() {
+  // Node's `parentPort` has its own message dispatch that is independent of the
+  // worker's global scope. Bun's native worker runtime dispatches parent messages
+  // onto the global scope (`self.onmessage` / `self.addEventListener('message')`),
+  // which matches web-worker semantics but not Node's — and it means emscripten-
+  // style code that does
+  //
+  //   parentPort.on('message', (m) => onmessage({ data: m }));
+  //   self.onmessage = handleMessage;
+  //
+  // sees every message delivered TWICE: once by the automatic `self.onmessage`
+  // dispatch and once by the explicit forwarding inside the `parentPort.on`
+  // listener. See https://github.com/oven-sh/bun/issues/29211.
+  //
+  // Fix: give `parentPort` its own `EventTarget`, re-dispatch incoming messages
+  // on it, and stop the native dispatch from reaching `self.onmessage` /
+  // `self.addEventListener('message', …)` so it matches Node semantics.
   const fake = Object.create(MessagePort.prototype);
+  const parentPortTarget = new EventTarget();
+
+  // Forwarders: installed lazily on `self` only while at least one user
+  // listener is registered on the parentPort. They intercept the native
+  // parent-message dispatch, stop immediate propagation (so `self.onmessage`
+  // and `self.addEventListener('message', …)` handlers on the global scope
+  // never see parent messages — matching Node), and re-dispatch on
+  // `parentPortTarget`. Installing a `message` listener on `self` keeps the
+  // event loop alive via `onDidChangeListenerImpl` in
+  // `BunWorkerGlobalScope.cpp`, so we only install while parentPort is
+  // actually being used — that way a worker that never touches `parentPort`
+  // exits cleanly when its module finishes executing.
+  //
+  // `onmessage` / `onmessageerror` are spec'd as implicit event listeners: the
+  // setter replaces at most one listener, firing in the order it was first
+  // assigned relative to other listeners on the same target.
+  let parentPortOnMessageWrapper: ((event: Event) => void) | null = null;
+  let parentPortOnMessageHandler: ((event: MessageEvent) => unknown) | null = null;
+  let parentPortOnMessageErrorWrapper: ((event: Event) => void) | null = null;
+  let parentPortOnMessageErrorHandler: ((event: MessageEvent) => unknown) | null = null;
+
+  // Separate counters per event type — installing a `message` listener on
+  // `self` keeps the event loop alive (via `onDidChangeListenerImpl` in
+  // `BunWorkerGlobalScope.cpp`, which only tracks `messageEvent`), so a
+  // worker that only registers a `messageerror` handler must NOT pull in the
+  // `message` forwarder.
+  let messageListenerCount = 0;
+  let messageErrorListenerCount = 0;
+  let messageForwarder: ((event: Event) => void) | null = null;
+  let messageErrorForwarder: ((event: Event) => void) | null = null;
+
+  const makeForwarder = (type: "message" | "messageerror") => (event: Event) => {
+    // Stop the native dispatch from reaching `self.onmessage` and any
+    // `self.addEventListener('message', …)` handlers — in Node parent
+    // messages are only visible through `parentPort`, not through the
+    // global scope.
+    event.stopImmediatePropagation();
+    const messageEvent = event as MessageEvent;
+    // Preserve `ports` so `worker.postMessage(data, [port])` still surfaces
+    // the transferred MessagePort(s) to `parentPort` listeners.
+    const nativePorts = messageEvent.ports;
+    const clone = new MessageEvent(type, {
+      data: messageEvent.data,
+      ports: nativePorts && nativePorts.length > 0 ? $Array.from(nativePorts) : undefined,
+    });
+    parentPortTarget.dispatchEvent(clone);
+  };
+
+  function installMessageForwarder() {
+    if (messageForwarder !== null) return;
+    messageForwarder = makeForwarder("message");
+    // Capture phase so we run before any user-installed bubbling listener
+    // on the global scope (if any).
+    self.addEventListener("message", messageForwarder, { capture: true });
+  }
+  function uninstallMessageForwarder() {
+    if (messageForwarder === null) return;
+    self.removeEventListener("message", messageForwarder, { capture: true } as any);
+    messageForwarder = null;
+  }
+  function installMessageErrorForwarder() {
+    if (messageErrorForwarder !== null) return;
+    messageErrorForwarder = makeForwarder("messageerror");
+    self.addEventListener("messageerror", messageErrorForwarder, { capture: true });
+  }
+  function uninstallMessageErrorForwarder() {
+    if (messageErrorForwarder === null) return;
+    self.removeEventListener("messageerror", messageErrorForwarder, { capture: true } as any);
+    messageErrorForwarder = null;
+  }
+
+  function acquireListener(type: "message" | "messageerror") {
+    if (type === "message") {
+      if (messageListenerCount++ === 0) installMessageForwarder();
+    } else {
+      if (messageErrorListenerCount++ === 0) installMessageErrorForwarder();
+    }
+  }
+
+  function releaseListener(type: "message" | "messageerror") {
+    if (type === "message") {
+      if (messageListenerCount > 0 && --messageListenerCount === 0) uninstallMessageForwarder();
+    } else {
+      if (messageErrorListenerCount > 0 && --messageErrorListenerCount === 0) uninstallMessageErrorForwarder();
+    }
+  }
+
+  // Wrap `addEventListener` / `removeEventListener` so we can track user
+  // listener lifetime on `parentPortTarget` and install / uninstall the
+  // forwarders on the global scope accordingly. Each (listener, type, capture)
+  // triple gets wrapped exactly once — duplicate adds are no-ops per the DOM
+  // spec — and the original listener object is the map key so that a
+  // `remove(type, original, {capture})` call finds the wrapped copy.
+  type TrackEntry = { wrapped: EventListener; once: boolean };
+  // `${type}:${capture ? 1 : 0}` — a listener registered with different
+  // (type, capture) combinations lives in separate slots, matching spec.
+  const trackedByListener = new WeakMap<object, Map<string, TrackEntry>>();
+
+  function listenerSlot(type: string, capture: boolean): string {
+    return capture ? type + ":1" : type + ":0";
+  }
+
+  function invokeListener(listener: EventListener | EventListenerObject, event: Event): void {
+    // DOM EventTarget accepts either a bare function or an object with a
+    // `handleEvent` method. Dispatch correctly for both forms.
+    if (typeof listener === "function") {
+      (listener as any).$call(fake, event);
+    } else if (listener !== null && typeof listener === "object" && typeof (listener as any).handleEvent === "function") {
+      (listener as any).handleEvent.$call(listener, event);
+    }
+  }
+
+  function parentPortAddEventListener(
+    type: string,
+    listener: EventListener | EventListenerObject | null,
+    options?: boolean | AddEventListenerOptions,
+  ): void {
+    if (listener === null || listener === undefined) return;
+    const capture = typeof options === "boolean" ? options : !!options?.capture;
+    const once = typeof options === "object" && options !== null && !!options.once;
+    // `AbortSignal` auto-removal is driven from the native EventTarget in C++,
+    // so it would bypass our JS `parentPortRemoveEventListener` wrapper and
+    // leak the event-loop refcount our capture forwarder holds on `self`.
+    // Strip the signal from the options we pass inward and re-implement abort
+    // ourselves via an abort listener that routes through the JS remove path.
+    const signal =
+      typeof options === "object" && options !== null ? ((options as AddEventListenerOptions).signal ?? null) : null;
+    if (signal && signal.aborted) return;
+    // Only `message` / `messageerror` events are dispatched on `self` by the
+    // native worker runtime — all other event types (`close`, `error`, …)
+    // live purely on `parentPortTarget` and don't need the capture forwarder
+    // at all.
+    const forwarderType: "message" | "messageerror" | null =
+      type === "message" ? "message" : type === "messageerror" ? "messageerror" : null;
+    const slot = listenerSlot(type, capture);
+    let bucket = trackedByListener.get(listener as object);
+    if (bucket?.$has(slot)) {
+      // Duplicate add — EventTarget already dedupes, so no-op.
+      return;
+    }
+    // Wrap so we can release the loop ref when the listener is removed,
+    // including the implicit removal done by `{ once: true }` after firing.
+    const wrapped: EventListener = function (event) {
+      if (once) {
+        const bucketNow = trackedByListener.get(listener as object);
+        if (bucketNow?.$get(slot) === entry) {
+          bucketNow.$delete(slot);
+          if (bucketNow.$size === 0) trackedByListener.delete(listener as object);
+          if (forwarderType !== null) releaseListener(forwarderType);
+        }
+      }
+      invokeListener(listener, event);
+    };
+    const entry: TrackEntry = { wrapped, once };
+    if (!bucket) {
+      bucket = new Map();
+      trackedByListener.set(listener as object, bucket);
+    }
+    bucket.$set(slot, entry);
+    const innerOptions: boolean | AddEventListenerOptions =
+      typeof options === "object" && options !== null ? { ...options, signal: undefined } : (options ?? false);
+    parentPortTarget.addEventListener(type, wrapped, innerOptions);
+    if (forwarderType !== null) acquireListener(forwarderType);
+    if (signal) {
+      signal.addEventListener(
+        "abort",
+        () => {
+          parentPortRemoveEventListener(type, listener, { capture });
+        },
+        { once: true },
+      );
+    }
+  }
+
+  function parentPortRemoveEventListener(
+    type: string,
+    listener: EventListener | EventListenerObject | null,
+    options?: boolean | EventListenerOptions,
+  ): void {
+    if (listener === null || listener === undefined) return;
+    const capture = typeof options === "boolean" ? options : !!options?.capture;
+    const bucket = trackedByListener.get(listener as object);
+    if (!bucket) return;
+    const slot = listenerSlot(type, capture);
+    const entry = bucket.$get(slot);
+    if (!entry) return;
+    bucket.$delete(slot);
+    if (bucket.$size === 0) trackedByListener.delete(listener as object);
+    parentPortTarget.removeEventListener(type, entry.wrapped, options);
+    if (type === "message") releaseListener("message");
+    else if (type === "messageerror") releaseListener("messageerror");
+  }
+
   Object.defineProperty(fake, "onmessage", {
     get() {
-      return self.onmessage;
+      return parentPortOnMessageHandler;
     },
     set(value) {
-      self.onmessage = value;
+      // Replace the previously-installed wrapper, if any.
+      if (parentPortOnMessageWrapper !== null) {
+        parentPortTarget.removeEventListener("message", parentPortOnMessageWrapper);
+        parentPortOnMessageWrapper = null;
+        releaseListener("message");
+      }
+      parentPortOnMessageHandler = typeof value === "function" ? value : null;
+      if (parentPortOnMessageHandler !== null) {
+        const handler = parentPortOnMessageHandler;
+        parentPortOnMessageWrapper = (event: Event) => {
+          try {
+            handler.$call(fake, event as MessageEvent);
+          } catch (err) {
+            queueMicrotask(() => {
+              throw err;
+            });
+          }
+        };
+        parentPortTarget.addEventListener("message", parentPortOnMessageWrapper);
+        acquireListener("message");
+      }
     },
   });
 
   Object.defineProperty(fake, "onmessageerror", {
     get() {
-      return self.onmessageerror;
+      return parentPortOnMessageErrorHandler;
     },
     set(value) {
-      self.onmessageerror = value;
+      if (parentPortOnMessageErrorWrapper !== null) {
+        parentPortTarget.removeEventListener("messageerror", parentPortOnMessageErrorWrapper);
+        parentPortOnMessageErrorWrapper = null;
+        releaseListener("messageerror");
+      }
+      parentPortOnMessageErrorHandler = typeof value === "function" ? value : null;
+      if (parentPortOnMessageErrorHandler !== null) {
+        const handler = parentPortOnMessageErrorHandler;
+        parentPortOnMessageErrorWrapper = (event: Event) => {
+          try {
+            handler.$call(fake, event as MessageEvent);
+          } catch (err) {
+            queueMicrotask(() => {
+              throw err;
+            });
+          }
+        };
+        parentPortTarget.addEventListener("messageerror", parentPortOnMessageErrorWrapper);
+        acquireListener("messageerror");
+      }
     },
   });
 
@@ -182,22 +430,35 @@
   });
 
   Object.defineProperty(fake, "addEventListener", {
-    value: self.addEventListener.bind(self),
+    value: parentPortAddEventListener,
   });
 
   Object.defineProperty(fake, "removeEventListener", {
-    value: self.removeEventListener.bind(self),
+    value: parentPortRemoveEventListener,
   });
 
-  Object.defineProperty(fake, "removeListener", {
-    value: self.removeEventListener.bind(self),
-    enumerable: false,
+  Object.defineProperty(fake, "dispatchEvent", {
+    value: parentPortTarget.dispatchEvent.bind(parentPortTarget),
   });
 
+  // `addListener`/`removeListener` must NOT bypass the `injectFakeEmitter`
+  // wrapping layer: `parentPort.on('message', fn)` wraps `fn` into a callback
+  // and stores `fn[wrappedListener] = callback`. A matching
+  // `parentPort.removeListener('message', fn)` has to resolve the wrapped
+  // callback before calling `removeEventListener`, otherwise it would try to
+  // remove `fn` itself (which was never registered) and silently leak the
+  // event-loop refcount held by our capture forwarder. `fake.on`/`fake.off`
+  // — inherited from `MessagePort.prototype` via `injectFakeEmitter` — do
+  // that resolution correctly, so we alias to them.
   Object.defineProperty(fake, "addListener", {
-    value: self.addEventListener.bind(self),
+    value: (MessagePort.prototype as any).on,
+    enumerable: false,
+  });
+
+  Object.defineProperty(fake, "removeListener", {
+    value: (MessagePort.prototype as any).off,
     enumerable: false,
   });
 
   return fake;
 }