FinalizationRegistry: keep the registry alive while it has registrations

src/jsc/bindings/JSCTaskScheduler.cpp

@@ -2,6 +2,10 @@
 #include <JavaScriptCore/VM.h>
 #include "JSCTaskScheduler.h"
 #include "BunClientData.h"
+#include <JavaScriptCore/JSFinalizationRegistry.h>
+#include <JavaScriptCore/StrongInlines.h>
+#include <JavaScriptCore/WeakMapImplInlines.h>
+#include <JavaScriptCore/JSCInlines.h>
 
 using Ticket = JSC::DeferredWorkTimer::Ticket;
 using Task = JSC::DeferredWorkTimer::Task;
@@ -96,6 +100,26 @@
         Bun__eventLoop__incrementRefConcurrently(bunVM, -1);
 }
 
+void JSCTaskScheduler::rootFinalizationRegistry(JSC::VM& vm, JSC::JSFinalizationRegistry* registry)
+{
+    ASSERT(vm.currentThreadIsHoldingAPILock());
+    auto result = m_rootedFinalizationRegistries.add(registry, JSC::Strong<JSC::JSObject>());
+    if (result.isNewEntry)
+        result.iterator->value.set(vm, registry);
+}
+
+void JSCTaskScheduler::unrootFinalizationRegistryIfDrained(JSC::JSFinalizationRegistry* registry)
+{
+    if (!m_rootedFinalizationRegistries.contains(registry))
+        return;
+    {
+        Locker cellLocker { registry->cellLock() };
+        if (registry->liveCount(cellLocker) || registry->deadCount(cellLocker))
+            return;
+    }
+    m_rootedFinalizationRegistries.remove(registry);
+}
+
 static void runPendingWork(void* bunVM, Bun::JSCTaskScheduler& scheduler, JSCDeferredWorkTask* job)
 {
     Locker<Lock> holder { scheduler.m_lock };
@@ -109,6 +133,8 @@
 
     if (pendingTicket && !pendingTicket->isCancelled()) {
         job->task(job->ticket.get());
+        if (auto* registry = dynamicDowncast<JSC::JSFinalizationRegistry>(job->ticket->target()))
+            scheduler.unrootFinalizationRegistryIfDrained(registry);
     }
 
     delete job;
@@ -127,9 +153,79 @@
 // has its enqueue visible to the drain; any that serializes after drops.
 extern "C" void Bun__JSCTaskScheduler__markShuttingDown(JSC::JSGlobalObject* globalObject)
 {
-    if (auto* clientData = WebCore::clientData(JSC::getVM(globalObject)))
+    if (auto* clientData = WebCore::clientData(JSC::getVM(globalObject))) {
+        clientData->deferredWorkTimer.m_rootedFinalizationRegistries.clear();
         clientData->deferredWorkTimer.markShuttingDown();
+    }
+}
+
+ALWAYS_INLINE static JSFinalizationRegistry* getFinalizationRegistry(VM& vm, JSGlobalObject* globalObject, JSValue value)
+{
+    auto scope = DECLARE_THROW_SCOPE(vm);
+    if (!value.isObject()) [[unlikely]] {
+        throwTypeError(globalObject, scope, "Called FinalizationRegistry function on non-object"_s);
+        return nullptr;
+    }
+    if (auto* registry = dynamicDowncast<JSFinalizationRegistry>(asObject(value))) [[likely]]
+        return registry;
+    throwTypeError(globalObject, scope, "Called FinalizationRegistry function on a non-FinalizationRegistry object"_s);
+    return nullptr;
+}
+
+JSC_DEFINE_HOST_FUNCTION(bunProtoFuncFinalizationRegistryRegister, (JSGlobalObject * globalObject, CallFrame* callFrame))
+{
+    VM& vm = globalObject->vm();
+    auto scope = DECLARE_THROW_SCOPE(vm);
+
+    auto* registry = getFinalizationRegistry(vm, globalObject, callFrame->thisValue());
+    RETURN_IF_EXCEPTION(scope, {});
+
+    JSValue target = callFrame->argument(0);
+    if (!canBeHeldWeakly(target)) [[unlikely]]
+        return throwVMTypeError(globalObject, scope, "register requires an object or a non-registered symbol as the target"_s);
+
+    JSValue holdings = callFrame->argument(1);
+    if (target == holdings) [[unlikely]]
+        return throwVMTypeError(globalObject, scope, "register expects the target object and the holdings parameter are not the same. Otherwise, the target can never be collected"_s);
+
+    JSValue unregisterToken = callFrame->argument(2);
+    if (!unregisterToken.isUndefined() && !canBeHeldWeakly(unregisterToken)) [[unlikely]]
+        return throwVMTypeError(globalObject, scope, "register requires an object or a non-registered symbol as the unregistration token"_s);
+
+    registry->registerTarget(vm, target.asCell(), holdings, unregisterToken);
+
+    if (auto* clientData = WebCore::clientData(vm))
+        clientData->deferredWorkTimer.rootFinalizationRegistry(vm, registry);
+    return encodedJSUndefined();
+}
+
+JSC_DEFINE_HOST_FUNCTION(bunProtoFuncFinalizationRegistryUnregister, (JSGlobalObject * globalObject, CallFrame* callFrame))
+{
+    VM& vm = globalObject->vm();
+    auto scope = DECLARE_THROW_SCOPE(vm);
+
+    auto* registry = getFinalizationRegistry(vm, globalObject, callFrame->thisValue());
+    RETURN_IF_EXCEPTION(scope, {});
+
+    JSValue token = callFrame->argument(0);
+    if (!canBeHeldWeakly(token)) [[unlikely]]
+        return throwVMTypeError(globalObject, scope, "unregister requires an object or a non-registered symbol as the unregistration token"_s);
+
+    bool result = registry->unregister(vm, token.asCell());
+    if (result) {
+        if (auto* clientData = WebCore::clientData(vm))
+            clientData->deferredWorkTimer.unrootFinalizationRegistryIfDrained(registry);
+    }
+    return JSValue::encode(jsBoolean(result));
+}
+
+void installFinalizationRegistryPrototypeHooks(JSC::JSGlobalObject* globalObject)
+{
+    VM& vm = globalObject->vm();
+    JSObject* prototype = globalObject->finalizationRegistryStructure()->storedPrototypeObject();
+    prototype->putDirectNativeFunction(vm, globalObject, Identifier::fromString(vm, "register"_s), 2, bunProtoFuncFinalizationRegistryRegister, ImplementationVisibility::Public, NoIntrinsic, static_cast<unsigned>(PropertyAttribute::DontEnum));
+    prototype->putDirectNativeFunction(vm, globalObject, Identifier::fromString(vm, "unregister"_s), 1, bunProtoFuncFinalizationRegistryUnregister, ImplementationVisibility::Public, NoIntrinsic, static_cast<unsigned>(PropertyAttribute::DontEnum));
 }
 
 // Reclaim a queued-but-never-dispatched job during shutdown. Called while the
 // JSC VM is still alive, so ~Ref<Ticket> and the captured Task lambda may

src/jsc/bindings/JSCTaskScheduler.h

@@ -5,6 +5,11 @@ class JSVMClientData;
 }
 
 #include <JavaScriptCore/DeferredWorkTimer.h>
+#include <JavaScriptCore/Strong.h>
+
+namespace JSC {
+class JSFinalizationRegistry;
+}
 
 namespace Bun {
 
@@ -20,6 +25,16 @@ class JSCTaskScheduler {
     static void onScheduleWorkSoon(WebCore::JSVMClientData* clientData, Ref<JSC::DeferredWorkTimer::Ticket>&& ticket, JSC::DeferredWorkTimer::Task&& task);
     static void onCancelPendingWork(WebCore::JSVMClientData* clientData, JSC::DeferredWorkTimer::Ticket& ticket);
 
+    // JavaScriptCore's async bytecode generator only preserves locals that are
+    // read after an `await`, so a `const fr = new FinalizationRegistry(...)`
+    // whose last use is `fr.register(...)` is collected at the next suspend
+    // point along with its pending registrations. V8 preserves every async
+    // local, so Node.js users never observe this. Root a registry on its first
+    // successful register() and release it once both its live and dead lists
+    // are empty so cleanup callbacks for already-registered targets still run.
+    void rootFinalizationRegistry(JSC::VM&, JSC::JSFinalizationRegistry*);
+    void unrootFinalizationRegistryIfDrained(JSC::JSFinalizationRegistry*);
+
     // Set once the owning VM's event loop has taken its last tick. After this,
     // onScheduleWorkSoon drops the task instead of enqueueing a ConcurrentTask
     // that can never be drained (~VM -> WaiterListManager::unregister reaches
@@ -37,6 +52,11 @@ class JSCTaskScheduler {
     bool m_isShuttingDown WTF_GUARDED_BY_LOCK(m_lock) { false };
     UncheckedKeyHashSet<Ref<JSC::DeferredWorkTimer::Ticket>> m_pendingTicketsKeepingEventLoopAlive;
     UncheckedKeyHashSet<Ref<JSC::DeferredWorkTimer::Ticket>> m_pendingTicketsOther;
+
+    // JS-thread only; see rootFinalizationRegistry above.
+    UncheckedKeyHashMap<JSC::JSCell*, JSC::Strong<JSC::JSObject>> m_rootedFinalizationRegistries;
 };
 
+void installFinalizationRegistryPrototypeHooks(JSC::JSGlobalObject*);
+
 }

src/jsc/bindings/ZigGlobalObject.cpp

@@ -3152,6 +3152,8 @@ void GlobalObject::addBuiltinGlobals(JSC::VM& vm)
 
     // ----- Extensions to Built-in objects -----
 
+    Bun::installFinalizationRegistryPrototypeHooks(this);
+
     JSC::JSObject* errorConstructor = this->errorConstructor();
     errorConstructor->putDirectNativeFunction(vm, this, JSC::Identifier::fromString(vm, "captureStackTrace"_s), 2, errorConstructorFuncCaptureStackTrace, ImplementationVisibility::Public, JSC::NoIntrinsic, PropertyAttribute::DontEnum | 0);
     errorConstructor->putDirectNativeFunction(vm, this, JSC::Identifier::fromString(vm, "appendStackTrace"_s), 2, errorConstructorFuncAppendStackTrace, ImplementationVisibility::Private, JSC::NoIntrinsic, PropertyAttribute::DontEnum | 0);

test/js/web/finalization-registry.test.ts

@@ -0,0 +1,207 @@
+import { test, expect, describe } from "bun:test";
+import { bunEnv, bunExe, tempDir } from "harness";
+
+// JavaScriptCore's async bytecode generator only preserves locals that are read
+// after an `await`, so a FinalizationRegistry held only by a module-level or
+// async-function `const` whose last use is `register()` becomes unreachable at
+// the next suspend point. Without an extra root the registry is swept before
+// its targets are observed as dead, and no cleanup callback ever fires (node
+// delivers them because V8 preserves every async-function local across await).
+
+async function run(source: string) {
+  // Module files, not `-e`: the eval wrapper's extra frames leave conservative
+  // stack roots that keep otherwise-dead cells alive for a few extra cycles.
+  using dir = tempDir("finalization-registry", { "entry.mjs": source });
+  await using proc = Bun.spawn({
+    cmd: [bunExe(), "entry.mjs"],
+    env: bunEnv,
+    cwd: String(dir),
+    stdout: "pipe",
+    stderr: "pipe",
+  });
+  const [stdout, stderr, exitCode] = await Promise.all([
+    proc.stdout.text(),
+    proc.stderr.text(),
+    proc.exited,
+  ]);
+  let json: unknown;
+  try {
+    json = JSON.parse(stdout.trim());
+  } catch {
+    json = undefined;
+  }
+  return { stdout, stderr, exitCode, json };
+}
+
+describe("FinalizationRegistry keeps itself alive while it has registrations", () => {
+  test.concurrent(
+    "cleanup callbacks fire when the registry local dies before its targets",
+    async () => {
+      const { stdout, stderr, exitCode, json } = await run(/* js */ `
+        const sleep = ms => new Promise(r => setTimeout(r, ms));
+        const K = 500;
+        let cleaned = 0;
+        const fr = new FinalizationRegistry(() => { cleaned++; });
+        const wrs = [];
+        for (let i = 0; i < K; i++) {
+          const o = { i, pad: Buffer.alloc(30, "p").toString() };
+          wrs.push(new WeakRef(o));
+          fr.register(o, i);
+        }
+        await sleep(300);
+        for (let r = 0; r < 30; r++) {
+          Bun.gc(true);
+          await sleep(10);
+          if (cleaned >= K) break;
+        }
+        const collected = wrs.filter(w => w.deref() === undefined).length;
+        console.log(JSON.stringify({ K, collected, cleaned }));
+      `);
+      expect(stderr).toBe("");
+      expect(stdout.trim()).not.toBe("");
+      const { K, collected, cleaned } = json as { K: number; collected: number; cleaned: number };
+      expect(collected).toBeGreaterThanOrEqual(K - 1);
+      // Without the rooting fix `cleaned` is 0 here; with it every collected
+      // target's callback runs.
+      expect(cleaned).toBe(collected);
+      expect(cleaned).toBeGreaterThanOrEqual(K - 1);
+      expect(exitCode).toBe(0);
+    },
+  );
+
+  test.concurrent("same inside an async function (not just module top level)", async () => {
+    const { stderr, exitCode, json } = await run(/* js */ `
+      const sleep = ms => new Promise(r => setTimeout(r, ms));
+      async function main() {
+        let cleaned = 0;
+        const fr = new FinalizationRegistry(() => { cleaned++; });
+        for (let i = 0; i < 200; i++) fr.register({ i }, i);
+        await 0;
+        for (let r = 0; r < 30; r++) {
+          Bun.gc(true);
+          await sleep(10);
+          if (cleaned >= 200) break;
+        }
+        console.log(JSON.stringify({ cleaned }));
+      }
+      await main();
+    `);
+    expect(stderr).toBe("");
+    expect((json as { cleaned: number }).cleaned).toBeGreaterThanOrEqual(199);
+    expect(exitCode).toBe(0);
+  });
+
+  test.concurrent("the registry is released once every registration is drained", async () => {
+    const { stderr, exitCode, json } = await run(/* js */ `
+      const sleep = ms => new Promise(r => setTimeout(r, ms));
+      let cleaned = 0;
+      globalThis.frw = undefined;
+      (function () {
+        const fr = new FinalizationRegistry(() => { cleaned++; });
+        globalThis.frw = new WeakRef(fr);
+        for (let i = 0; i < 100; i++) fr.register({ i }, i);
+      })();
+      for (let r = 0; r < 30; r++) {
+        Bun.gc(true);
+        await sleep(10);
+        if (cleaned >= 100) break;
+      }
+      for (let r = 0; r < 30; r++) {
+        Bun.gc(true);
+        await sleep(10);
+        if (globalThis.frw.deref() === undefined) break;
+      }
+      console.log(JSON.stringify({ cleaned, drained: globalThis.frw.deref() === undefined }));
+    `);
+    expect(stderr).toBe("");
+    const { cleaned, drained } = json as { cleaned: number; drained: boolean };
+    expect(cleaned).toBe(100);
+    // Released once drained; otherwise the Strong root would leak it forever.
+    expect(drained).toBe(true);
+    expect(exitCode).toBe(0);
+  });
+
+  // Conservative stack scanning can pin any one cell for a few cycles; use a
+  // batch of registries so a stray stack word can only shadow the count, never
+  // the invariant that the Strong root is released.
+  test.concurrent("unregister() that drains every entry releases the root", async () => {
+    const { stderr, exitCode, json } = await run(/* js */ `
+      const sleep = ms => new Promise(r => setTimeout(r, ms));
+      const wrs = [];
+      (function () {
+        const tok = {};
+        for (let n = 0; n < 200; n++) {
+          const fr = new FinalizationRegistry(() => {});
+          wrs.push(new WeakRef(fr));
+          for (let i = 0; i < 4; i++) fr.register(globalThis, i, tok);
+          fr.unregister(tok);
+        }
+      })();
+      for (let r = 0; r < 30; r++) {
+        Bun.gc(true);
+        await sleep(5);
+        if (wrs.every(w => w.deref() === undefined)) break;
+      }
+      const alive = wrs.filter(w => w.deref() !== undefined).length;
+      console.log(JSON.stringify({ alive, total: wrs.length }));
+    `);
+    expect(stderr).toBe("");
+    const { alive, total } = json as { alive: number; total: number };
+    expect(total).toBe(200);
+    expect(alive).toBeLessThanOrEqual(2);
+    expect(exitCode).toBe(0);
+  });
+
+  test.concurrent("a registry that never registers stays collectable", async () => {
+    const { stderr, exitCode, json } = await run(/* js */ `
+      const sleep = ms => new Promise(r => setTimeout(r, ms));
+      const wrs = [];
+      (function () {
+        for (let n = 0; n < 200; n++) {
+          const fr = new FinalizationRegistry(() => {});
+          wrs.push(new WeakRef(fr));
+        }
+      })();
+      for (let r = 0; r < 30; r++) {
+        Bun.gc(true);
+        await sleep(5);
+        if (wrs.every(w => w.deref() === undefined)) break;
+      }
+      const alive = wrs.filter(w => w.deref() !== undefined).length;
+      console.log(JSON.stringify({ alive, total: wrs.length }));
+    `);
+    expect(stderr).toBe("");
+    const { alive, total } = json as { alive: number; total: number };
+    expect(total).toBe(200);
+    expect(alive).toBeLessThanOrEqual(2);
+    expect(exitCode).toBe(0);
+  });
+
+  test.concurrent("register/unregister argument validation is unchanged", async () => {
+    const { stderr, exitCode, json } = await run(/* js */ `
+      const fr = new FinalizationRegistry(() => {});
+      const errors = [];
+      const catchType = fn => { try { fn(); errors.push(null); } catch (e) { errors.push(e?.constructor?.name); } };
+      catchType(() => fr.register(42, "x"));
+      catchType(() => fr.register({}, "x", 42));
+      const obj = {};
+      catchType(() => fr.register(obj, obj));
+      catchType(() => fr.unregister(42));
+      catchType(() => FinalizationRegistry.prototype.register.call({}, {}, "x"));
+      const tok = {};
+      fr.register({}, "x", tok);
+      const ok = fr.unregister(tok);
+      const okTwice = fr.unregister(tok);
+      console.log(JSON.stringify({ errors, ok, okTwice, len: fr.register.length, ulen: fr.unregister.length }));
+    `);
+    expect(stderr).toBe("");
+    expect(json).toEqual({
+      errors: ["TypeError", "TypeError", "TypeError", "TypeError", "TypeError"],
+      ok: true,
+      okTwice: false,
+      len: 2,
+      ulen: 1,
+    });
+    expect(exitCode).toBe(0);
+  });
+});