From b53c25e5f819e9b93cfd456bc07fc15bdf76cee6 Mon Sep 17 00:00:00 2001
From: Jarred Sumner <jarred@jarredsumner.com>
Date: Fri, 23 Aug 2024 21:09:56 -0700
Subject: [PATCH] Redo napi cleanup hooks (#13487)

---
 src/bun.js/javascript.zig | 10 +++++-----
 src/bun.js/rare_data.zig  | 18 +++++-------------
 src/napi/napi.zig         | 32 ++++++++++++++------------------
 3 files changed, 24 insertions(+), 36 deletions(-)

diff --git a/src/bun.js/javascript.zig b/src/bun.js/javascript.zig
index 615932bb64..64a476e67c 100644
--- a/src/bun.js/javascript.zig
+++ b/src/bun.js/javascript.zig
@@ -1304,11 +1304,11 @@ pub const VirtualMachine = struct {
         this.exit_handler.dispatchOnExit();
 
         const rare_data = this.rare_data orelse return;
-        var hook = rare_data.cleanup_hook orelse return;
-        hook.execute();
-        while (hook.next) |next| {
-            next.execute();
-            hook = next;
+        var hooks = rare_data.cleanup_hooks;
+        defer if (!is_main_thread_vm) hooks.clearAndFree(bun.default_allocator);
+        rare_data.cleanup_hooks = .{};
+        for (hooks.items) |hook| {
+            hook.execute();
         }
     }
 
diff --git a/src/bun.js/rare_data.zig b/src/bun.js/rare_data.zig
index 6451180184..45ebac3282 100644
--- a/src/bun.js/rare_data.zig
+++ b/src/bun.js/rare_data.zig
@@ -31,8 +31,7 @@ hot_map: ?HotMap = null,
 
 // TODO: make this per JSGlobalObject instead of global
 // This does not handle ShadowRealm correctly!
-tail_cleanup_hook: ?*CleanupHook = null,
-cleanup_hook: ?*CleanupHook = null,
+cleanup_hooks: std.ArrayListUnmanaged(CleanupHook) = .{},
 
 file_polls_: ?*Async.FilePoll.Store = null,
 
@@ -220,7 +219,6 @@ pub const EntropyCache = struct {
 };
 
 pub const CleanupHook = struct {
-    next: ?*CleanupHook = null,
     ctx: ?*anyopaque,
     func: Function,
     globalThis: *JSC.JSGlobalObject,
@@ -233,13 +231,12 @@ pub const CleanupHook = struct {
         self.func(self.ctx);
     }
 
-    pub fn from(
+    pub fn init(
         globalThis: *JSC.JSGlobalObject,
         ctx: ?*anyopaque,
         func: CleanupHook.Function,
     ) CleanupHook {
         return .{
-            .next = null,
             .ctx = ctx,
             .func = func,
             .globalThis = globalThis,
@@ -255,14 +252,7 @@ pub fn pushCleanupHook(
     ctx: ?*anyopaque,
     func: CleanupHook.Function,
 ) void {
-    const hook = JSC.VirtualMachine.get().allocator.create(CleanupHook) catch unreachable;
-    hook.* = CleanupHook.from(globalThis, ctx, func);
-    if (this.cleanup_hook == null) {
-        this.cleanup_hook = hook;
-        this.tail_cleanup_hook = hook;
-    } else {
-        this.cleanup_hook.?.next = hook;
-    }
+    this.cleanup_hooks.append(bun.default_allocator, CleanupHook.init(globalThis, ctx, func)) catch bun.outOfMemory();
 }
 
 pub fn boringEngine(rare: *RareData) *BoringSSL.ENGINE {
@@ -402,4 +392,6 @@ pub fn deinit(this: *RareData) void {
     if (this.boring_ssl_engine) |engine| {
         _ = bun.BoringSSL.ENGINE_free(engine);
     }
+
+    this.cleanup_hooks.clearAndFree(bun.default_allocator);
 }
diff --git a/src/napi/napi.zig b/src/napi/napi.zig
index c6f0d58ede..3422727ef3 100644
--- a/src/napi/napi.zig
+++ b/src/napi/napi.zig
@@ -1335,28 +1335,24 @@ pub export fn napi_add_env_cleanup_hook(env: napi_env, fun: ?*const fn (?*anyopa
 }
 pub export fn napi_remove_env_cleanup_hook(env: napi_env, fun: ?*const fn (?*anyopaque) callconv(.C) void, arg: ?*anyopaque) napi_status {
     log("napi_remove_env_cleanup_hook", .{});
-    if (env.bunVM().rare_data == null or fun == null)
+
+    // Avoid looking up env.bunVM().
+    if (bun.Global.isExiting()) {
+        return .ok;
+    }
+
+    const vm = JSC.VirtualMachine.get();
+
+    if (vm.rare_data == null or fun == null or vm.isShuttingDown())
         return .ok;
 
-    var rare_data = env.bunVM().rare_data.?;
-    var hook = rare_data.cleanup_hook orelse return .ok;
-    const cmp = JSC.RareData.CleanupHook.from(env, arg, fun.?);
-    if (hook.eql(cmp)) {
-        env.bunVM().allocator.destroy(hook);
-        rare_data.cleanup_hook = null;
-        rare_data.tail_cleanup_hook = null;
-    }
-    while (hook.next) |current| {
+    var rare_data = vm.rare_data.?;
+    const cmp = JSC.RareData.CleanupHook.init(env, arg, fun.?);
+    for (rare_data.cleanup_hooks.items, 0..) |*hook, i| {
         if (hook.eql(cmp)) {
-            if (current.next) |next| {
-                hook.next = next;
-            } else {
-                hook.next = null;
-            }
-            env.bunVM().allocator.destroy(current);
-            return .ok;
+            _ = rare_data.cleanup_hooks.orderedRemove(i);
+            break;
         }
-        hook = current;
     }
 
     return .ok;