Release a JSC deferred-work keep-alive on the event loop that took it

src/jsc/bindings/JSCTaskScheduler.cpp

@@ -36,20 +36,25 @@ class JSCDeferredWorkTask {
     WTF_MAKE_TZONE_ALLOCATED(JSCDeferredWorkTask);
 };
 
-// Drop `ticket` from whichever pending set holds it. Caller holds m_lock; the
-// event-loop ref is balanced after the caller releases the lock.
-static bool dropPendingTicketLocked(Bun::JSCTaskScheduler& scheduler, Ticket* ticket) WTF_REQUIRES_LOCK(scheduler.m_lock)
+// Drop `ticket` from whichever pending set holds it. Caller holds m_lock. Returns the loop
+// whose keep-alive the ticket held, if it held one; the caller releases it on that loop after
+// it unlocks.
+static std::optional<BunLoopKind> dropPendingTicketLocked(Bun::JSCTaskScheduler& scheduler, Ticket* ticket) WTF_REQUIRES_LOCK(scheduler.m_lock)
 {
-    bool isKeepingEventLoopAlive = scheduler.m_pendingTicketsKeepingEventLoopAlive.removeIf([ticket](auto& pendingTicket) {
-        return pendingTicket.key.ptr() == ticket;
+    std::optional<BunLoopKind> keptAlive;
+    scheduler.m_pendingTicketsKeepingEventLoopAlive.removeIf([&](auto& pendingTicket) {
+        if (pendingTicket.key.ptr() != ticket)
+            return false;
+        keptAlive = pendingTicket.value.loopKind;
+        return true;
     });
     // -- At this point, ticket may be an invalid pointer.
-    if (!isKeepingEventLoopAlive) {
+    if (!keptAlive) {
         scheduler.m_pendingTicketsOther.removeIf([ticket](auto& pendingTicket) {
             return pendingTicket.key.ptr() == ticket;
         });
     }
-    return isKeepingEventLoopAlive;
+    return keptAlive;
 }
 
 void JSCTaskScheduler::onAddPendingWork(WebCore::JSVMClientData* clientData, Ref<Ticket>&& ticket, JSC::DeferredWorkTimer::WorkType kind)
@@ -86,10 +91,10 @@ void JSCTaskScheduler::onScheduleWorkSoon(WebCore::JSVMClientData* clientData, R
         // collectNow -> JSFinalizationRegistry::finalizeUnconditionally. Balance
         // onAddPendingWork so the ticket-set entry and event-loop ref are released.
         if (scheduler.m_isShuttingDown) [[unlikely]] {
-            bool wasKeepingAlive = dropPendingTicketLocked(scheduler, ticket.ptr());
+            auto keptAlive = dropPendingTicketLocked(scheduler, ticket.ptr());
             holder.unlockEarly();
-            if (wasKeepingAlive)
-                Bun__VmHandle__refKeepAlive(clientData->vmHandle, BunLoopKind::Regular, -1);
+            if (keptAlive)
+                Bun__VmHandle__refKeepAlive(clientData->vmHandle, *keptAlive, -1);
             return;
         }
         auto it = scheduler.m_pendingTicketsKeepingEventLoopAlive.find(ticket.ptr());
@@ -108,10 +113,10 @@ void JSCTaskScheduler::onCancelPendingWork(WebCore::JSVMClientData* clientData,
     auto& scheduler = clientData->deferredWorkTimer;
 
     Locker<Lock> holder { scheduler.m_lock };
-    bool wasKeepingAlive = dropPendingTicketLocked(scheduler, &ticket);
+    auto keptAlive = dropPendingTicketLocked(scheduler, &ticket);
     holder.unlockEarly();
-    if (wasKeepingAlive)
-        Bun__VmHandle__refKeepAlive(vmHandle, BunLoopKind::Regular, -1);
+    if (keptAlive)
+        Bun__VmHandle__refKeepAlive(vmHandle, *keptAlive, -1);
 }
 
 static void runPendingWork(const ::BunVmHandleRef* vmHandle, Bun::JSCTaskScheduler& scheduler, JSCDeferredWorkTask* job)
@@ -121,9 +126,10 @@ static void runPendingWork(const ::BunVmHandleRef* vmHandle, Bun::JSCTaskSchedul
     uint32_t graphContext = 0;
     if (auto it = scheduler.m_pendingTicketsKeepingEventLoopAlive.find(job->ticket.ptr()); it != scheduler.m_pendingTicketsKeepingEventLoopAlive.end()) {
         graphContext = it->value.graphContext;
+        BunLoopKind loopKind = it->value.loopKind;
         scheduler.m_pendingTicketsKeepingEventLoopAlive.remove(it);
         wasPending = true;
-        Bun__VmHandle__refKeepAlive(vmHandle, BunLoopKind::Regular, -1);
+        Bun__VmHandle__refKeepAlive(vmHandle, loopKind, -1);
     } else if (auto it = scheduler.m_pendingTicketsOther.find(job->ticket.ptr()); it != scheduler.m_pendingTicketsOther.end()) {
         graphContext = it->value.graphContext;
         scheduler.m_pendingTicketsOther.remove(it);
@@ -182,10 +188,10 @@ extern "C" void Bun__deleteDeferredWorkTask(Bun::JSCDeferredWorkTask* job)
     if (auto* clientData = WebCore::clientData(job->vm())) {
         auto& scheduler = clientData->deferredWorkTimer;
         Locker<Lock> holder { scheduler.m_lock };
-        bool wasKeepingAlive = dropPendingTicketLocked(scheduler, job->ticket.ptr());
+        auto keptAlive = dropPendingTicketLocked(scheduler, job->ticket.ptr());
         holder.unlockEarly();
-        if (wasKeepingAlive)
-            Bun__VmHandle__refKeepAlive(clientData->vmHandle, BunLoopKind::Regular, -1);
+        if (keptAlive)
+            Bun__VmHandle__refKeepAlive(clientData->vmHandle, *keptAlive, -1);
     }
     delete job;
 }

src/jsc/bindings/JSCTaskScheduler.h

@@ -35,7 +35,7 @@ class JSCTaskScheduler {
 public:
     // What was current when JSC registered the work.
     struct PendingWork {
-        // Its completion is posted to this loop.
+        // Its completion is posted to this loop, and the keep-alive it took there is released on it.
         BunLoopKind loopKind { BunLoopKind::Regular };
         // The identifier of the Bun.ModuleGraph context whose script asked for the work, or 0 for
         // the realm's own: the completion of a stopped one is dropped.

test/bundler/transpiler/macro-test.test.ts

@@ -313,8 +313,9 @@ test("a Response or Blob returned from a macro is classified by its MIME essence
 // loop was current when their work started: what the macro started goes to the macro loop (or the wait
 // hangs), what the program started stays on the regular loop (or program callbacks run mid-transpile),
 // and whatever a macro started but did not await is adopted by the regular loop once the macro returns
-// (or it is stranded and its keep-alive holds the process open). These run the macro in the main VM:
-// the entry file's macros, or a module require()d so it transpiles on the main thread.
+// (or it is stranded and its keep-alive holds the process open). Unless a test names another VM, these
+// run the macro in the main VM: the entry file's macros, or a module require()d so it transpiles on
+// the main thread.
 describe("event loop routing around macros", () => {
   async function run(files: Record<string, string>, env: Record<string, string> = {}) {
     using dir = tempDir("macro-loops", files);
@@ -451,6 +452,60 @@ describe("event loop routing around macros", () => {
     expect({ lines, stderr }).toEqual({ lines: ["1 chained"], stderr: "" });
     expect(exitCode).toBe(0);
   });
+
+  // JSC takes a keep-alive for a WebAssembly.compile() on the loop that is current, the macro loop here,
+  // and has to release it on that same loop. A macro VM on a transpiler or bundler thread never ticks
+  // its regular loop, so a release parked there is never applied and the thread's loop stays active.
+  const keepAliveMacro = [
+    `import { getEventLoopStats } from "bun:internal-for-testing";`,
+    `let before = 0;`,
+    `export function start() {`,
+    `  before = getEventLoopStats().numPolls;`,
+    `  return 0;`,
+    `}`,
+    `export async function compile() {`,
+    `  await WebAssembly.compile(new Uint8Array([0, 0x61, 0x73, 0x6d, 1, 0, 0, 0]));`,
+    `  return 0;`,
+    `}`,
+    `export function leaked() {`,
+    `  const thread = Bun.isMainThread ? "the main thread" : "another thread";`,
+    `  return (getEventLoopStats().numPolls - before) + " on " + thread;`,
+    `}`,
+  ].join("\n");
+  const callsKeepAliveMacro = [
+    `import { start, compile, leaked } from "./m.ts" with { type: "macro" };`,
+    `start();`,
+    `compile();`,
+    `console.log("leaked", leaked());`,
+  ].join("\n");
+  const macroVMs: [name: string, files: Record<string, string>, line: string][] = [
+    ["the main VM", { "index.ts": callsKeepAliveMacro }, "leaked 0 on the main thread"],
+    [
+      "a transpiler thread's VM",
+      { "lib.ts": callsKeepAliveMacro, "index.ts": `import "./lib.ts";\n` },
+      "leaked 0 on another thread",
+    ],
+    [
+      "a Bun.build() worker's VM",
+      {
+        "entry.ts": callsKeepAliveMacro,
+        "index.ts": [
+          `const result = await Bun.build({ entrypoints: ["./entry.ts"] });`,
+          `console.log((await result.outputs[0].text()).trim().split("\\n").pop());`,
+        ].join("\n"),
+      },
+      `console.log("leaked", "0 on another thread");`,
+    ],
+  ];
+
+  test.concurrent.each(macroVMs)(
+    "a WebAssembly.compile() that a macro awaits in %s releases its event loop keep-alive",
+    async (_name, files, line) => {
+      const { lines, stderr, exitCode } = await run({ "m.ts": keepAliveMacro, ...files });
+      expect({ lines, stderr }).toEqual({ lines: [line], stderr: "" });
+      expect(exitCode).toBe(0);
+    },
+  );
 });
 
 // A module that is not the entry point is transpiled on a worker thread, where no VM exists yet. The