Run FinalizationRegistry cleanup callbacks

runtime/engine.cpp

@@ -601,14 +601,21 @@ bool Engine::eval_toplevel(JS::SourceText<mozilla::Utf8Unit> &source, std::strin
   // TLA rejections during pre-initialization are treated as top-level exceptions.
   // TLA may remain unresolved, in which case it will continue tasks at runtime.
   // Rejections after pre-intialization remain unhandled rejections for now.
-  if (tla_promise.isObject()) {
+  auto tla_rejected = [&]() {
+    if (!tla_promise.isObject()) {
+      return false;
+    }
     RootedObject promise_obj(cx(), &tla_promise.toObject());
-    JS::PromiseState state = JS::GetPromiseState(promise_obj);
-    if (state == JS::PromiseState::Rejected) {
-      RootedValue err(cx(), JS::GetPromiseResult(promise_obj));
-      JS_SetPendingException(cx(), err);
+    if (JS::GetPromiseState(promise_obj) != JS::PromiseState::Rejected) {
       return false;
     }
+    RootedValue err(cx(), JS::GetPromiseResult(promise_obj));
+    JS_SetPendingException(cx(), err);
+    return true;
+  };
+
+  if (tla_rejected()) {
+    return false;
   }
 
   // Report any promise rejections that weren't handled before snapshotting.
@@ -638,6 +645,18 @@ bool Engine::eval_toplevel(JS::SourceText<mozilla::Utf8Unit> &source, std::strin
   if (state() == EngineState::ScriptPreInitializing) {
     JS::PrepareForFullGC(cx());
     JS::NonIncrementalGC(cx(), JS::GCOptions::Normal, JS::GCReason::API);
+
+    // The collection may have queued FinalizationRegistry cleanups. Run them now, so they are
+    // not baked into the snapshot and run on its first resumption.
+    if (!core::EventLoop::drain_finalization_registry_cleanups(this)) {
+      return false;
+    }
+
+    // A cleanup, or a microtask it queued, may have rejected the TLA promise. Check again, so
+    // that a failing module is not snapshotted.
+    if (tla_rejected()) {
+      return false;
+    }
   }
 
   // Ignore the first GC, but then print all others, because ideally GCs

runtime/event_loop.cpp

@@ -3,6 +3,8 @@
 #include "host_api.h"
 #include "jsapi.h"
 #include "jsfriendapi.h"
+#include "js/GCVector.h"
+#include "js/Utility.h"
 
 #include <iostream>
 #include <print>
@@ -22,6 +24,48 @@ struct TaskQueue {
 
 static PersistentRooted<TaskQueue> queue;
 
+// SpiderMonkey hands the embedding one cleanup function per FinalizationRegistry that has
+// records of dead targets, and leaves calling it to the embedding (HTML's
+// HostEnqueueFinalizationRegistryCleanupJob). The functions are queued here and run by the event
+// loop, as a task after the microtask checkpoint.
+using CleanupFunctions = JS::GCVector<JSFunction *, 0, js::SystemAllocPolicy>;
+static PersistentRooted<CleanupFunctions> finalization_registry_cleanups;
+
+// Called during GC, so this must not trigger another collection: it only appends to the queue.
+// An append can only fail on allocation failure, which is fatal, as in SpiderMonkey's shell.
+static void queue_finalization_registry_cleanup(JSFunction *do_cleanup,
+                                                JSObject * /*incumbent_global*/,
+                                                void * /*data*/) {
+  js::AutoEnterOOMUnsafeRegion oom_unsafe;
+  if (!finalization_registry_cleanups.get().append(do_cleanup)) {
+    oom_unsafe.crash("queue_finalization_registry_cleanup");
+  }
+}
+
+// Calls every queued cleanup function and returns whether any ran. An exception thrown by a
+// callback is reported and cleared, and does not fail the event loop, as in a browser.
+static bool run_finalization_registry_cleanups(api::Engine *engine) {
+  if (finalization_registry_cleanups.get().empty()) {
+    return false;
+  }
+  JSContext *cx = engine->cx();
+  JS::Rooted<CleanupFunctions> cleanups(cx);
+  std::swap(cleanups.get(), finalization_registry_cleanups.get());
+  JS::RootedFunction cleanup(cx);
+  JS::RootedValue rval(cx);
+  for (JSFunction *fn : cleanups) {
+    cleanup = fn;
+    JS::RootedObject cleanup_obj(cx, JS_GetFunctionObject(cleanup));
+    JS::ExposeObjectToActiveJS(cleanup_obj);
+    JSAutoRealm ar(cx, cleanup_obj);
+    if (!JS_CallFunction(cx, nullptr, cleanup, JS::HandleValueArray::empty(), &rval)) {
+      api::Engine::dump_pending_exception("running a FinalizationRegistry callback");
+      JS_ClearPendingException(cx);
+    }
+  }
+  return true;
+}
+
 namespace core {
 
 void EventLoop::queue_async_task(const RefPtr<api::AsyncTask>& task) {
@@ -70,6 +114,13 @@ bool EventLoop::run_event_loop(api::Engine *engine, double total_compute) {
       exit_event_loop();
       return false;
     }
+
+    // Once the microtasks are done, run the FinalizationRegistry cleanups that a collection
+    // queued, then go round again for the microtasks they queued.
+    if (run_finalization_registry_cleanups(engine)) {
+      continue;
+    }
+
     // if there is no interest in the event loop at all, just run one tick
     if (interest_complete()) {
       exit_event_loop();
@@ -97,6 +148,24 @@ bool EventLoop::run_event_loop(api::Engine *engine, double total_compute) {
   }
 }
 
-void EventLoop::init(JSContext *cx) { queue.init(cx); }
+bool EventLoop::drain_finalization_registry_cleanups(api::Engine *engine) {
+  JSContext *cx = engine->cx();
+  for (unsigned round = 0; round < MAX_CLEANUP_DRAIN_ROUNDS; round++) {
+    if (!run_finalization_registry_cleanups(engine)) {
+      return true;
+    }
+    js::RunJobs(cx);
+    if (JS_IsExceptionPending(cx)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+void EventLoop::init(JSContext *cx) {
+  queue.init(cx);
+  finalization_registry_cleanups.init(cx);
+  JS::SetHostCleanupFinalizationRegistryCallback(cx, queue_finalization_registry_cleanup, nullptr);
+}
 
 } // namespace core

runtime/event_loop.h

@@ -24,6 +24,17 @@ class EventLoop {
    */
   static bool run_event_loop(api::Engine *engine, double total_compute);
 
+  /**
+   * Run the FinalizationRegistry cleanups that collections have queued, and the microtasks they
+   * queue, until none are left. Used before a Wizer snapshot is taken, so that the snapshot holds
+   * an empty queue. A cleanup that triggers a collection that queues more is run in a further
+   * round, up to MAX_CLEANUP_DRAIN_ROUNDS; anything still queued then stays queued, and runs in
+   * the first event loop after the snapshot is resumed. Returns false if an exception is pending.
+   */
+  static bool drain_finalization_registry_cleanups(api::Engine *engine);
+
+  static constexpr unsigned MAX_CLEANUP_DRAIN_ROUNDS = 16;
+
   static void incr_event_loop_interest();
   static void decr_event_loop_interest();
 

tests/e2e/finalization-registry-snapshot/expect_serve_body.txt

@@ -0,0 +1 @@
+finalized 1

tests/e2e/finalization-registry-snapshot/finalization-registry-snapshot.js

@@ -0,0 +1,16 @@
+// The collection that runs before the Wizer snapshot is taken finds the target registered below
+// dead. Its cleanup must have run before the snapshot, not be baked into it and run after it is
+// resumed: the handler answers synchronously, before the first event loop of the resumed
+// instance, so it sees the count only if the cleanup ran during pre-initialization.
+let finalized = 0;
+const registry = new FinalizationRegistry(() => {
+  finalized += 1;
+});
+
+// In a function, so that no top-level frame still refers to the target.
+function register() {
+  registry.register({}, 0);
+}
+register();
+
+addEventListener("fetch", (event) => event.respondWith(new Response(`finalized ${finalized}`)));

tests/e2e/finalization-registry-tla-err/expect_wizer_stderr.txt

@@ -0,0 +1,2 @@
+Exception while evaluating top-level script
+e2e/finalization-registry-tla-err/finalization-registry-tla-err.js:6:10 Error: rejected by a FinalizationRegistry cleanup

tests/e2e/finalization-registry-tla-err/finalization-registry-tla-err.js

@@ -0,0 +1,15 @@
+// The collection that runs before the Wizer snapshot is taken finds the target registered below
+// dead, and its cleanup rejects the pending top-level await. Pre-initialization must fail, as it
+// does for any TLA rejection, rather than snapshot a failed module.
+let reject;
+const registry = new FinalizationRegistry(() => {
+  reject(new Error('rejected by a FinalizationRegistry cleanup'));
+});
+
+// In a function, so that no top-level frame still refers to the target.
+function register() {
+  registry.register({}, 0);
+}
+register();
+
+await new Promise((_, _reject) => void (reject = _reject));

tests/e2e/finalization-registry/expect_serve_body.txt

@@ -0,0 +1 @@
+finalized

tests/e2e/finalization-registry/finalization-registry.js

@@ -0,0 +1,24 @@
+// A FinalizationRegistry callback runs once a collection has found its target dead. The
+// collection is started by allocation, and the callback runs in the event loop, after the
+// microtask checkpoint that follows it.
+let finalized = 0;
+const registry = new FinalizationRegistry(() => {
+  finalized += 1;
+});
+
+async function handle() {
+  registry.register({}, 0);
+
+  // Allocate until a collection has run and its cleanup has been called, giving up after a
+  // bounded number of attempts.
+  for (let i = 0; i < 1000 && finalized === 0; i++) {
+    const buffer = new Uint8Array(2 << 20);
+    buffer[0] = 1;
+    // Wait for a task, so that the cleanups queued by a collection can run.
+    await new Promise((resolve) => setTimeout(resolve, 0));
+  }
+
+  return new Response(finalized > 0 ? "finalized" : "not finalized");
+}
+
+addEventListener("fetch", (event) => event.respondWith(handle()));

tests/tests.cmake

@@ -61,6 +61,9 @@ test_e2e(teed-stream-as-outgoing-body)
 test_e2e(init-script)
 test_e2e(no-init-location)
 test_e2e(init-location)
+test_e2e(finalization-registry)
+test_e2e(finalization-registry-snapshot)
+test_e2e(finalization-registry-tla-err)
 
 integration_tests(
     blob