cleanup hook

4 files changed, 111 insertions, 5 deletions

diff --git a/src/bun_js.zig b/src/bun_js.zig
index 7fc2f43af..69fe006c6 100644
--- a/src/bun_js.zig
+++ b/src/bun_js.zig
@@ -159,6 +159,8 @@ pub const Run = struct {
             }
         }
 
+        this.vm.onExit();
+
         Global.exit(0);
     }
 };
diff --git a/src/javascript/jsc/javascript.zig b/src/javascript/jsc/javascript.zig
index 06f0c47a1..f6147a1dd 100644
--- a/src/javascript/jsc/javascript.zig
+++ b/src/javascript/jsc/javascript.zig
@@ -584,6 +584,16 @@ pub const VirtualMachine = struct {
         return this.event_loop;
     }
 
+    pub fn onExit(this: *VirtualMachine) void {
+        var 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;
+        }
+    }
+
     pub const EventLoop = struct {
         ready_tasks_count: std.atomic.Atomic(u32) = std.atomic.Atomic(u32).init(0),
         pending_tasks_count: std.atomic.Atomic(u32) = std.atomic.Atomic(u32).init(0),
diff --git a/src/javascript/jsc/rare_data.zig b/src/javascript/jsc/rare_data.zig
index b6559826c..90b8861d4 100644
--- a/src/javascript/jsc/rare_data.zig
+++ b/src/javascript/jsc/rare_data.zig
@@ -14,6 +14,55 @@ stderr_store: ?*Blob.Store = null,
 stdin_store: ?*Blob.Store = null,
 stdout_store: ?*Blob.Store = null,
 
+// TODO: make this per JSGlobalObject instead of global
+// This does not handle ShadowRealm correctly!
+tail_cleanup_hook: ?*CleanupHook = null,
+cleanup_hook: ?*CleanupHook = null,
+
+pub const CleanupHook = struct {
+    next: ?*CleanupHook = null,
+    ctx: ?*anyopaque,
+    func: Function,
+    globalThis: *JSC.JSGlobalObject,
+
+    pub fn eql(self: CleanupHook, other: CleanupHook) bool {
+        return self.ctx == other.ctx and self.func == other.func and self.globalThis == other.globalThis;
+    }
+
+    pub fn from(
+        globalThis: *JSC.JSGlobalObject,
+        ctx: ?*anyopaque,
+        func: CleanupHook.Function,
+    ) CleanupHook {
+        return .{
+            .next = null,
+            .ctx = ctx,
+            .func = func,
+            .globalThis = globalThis,
+        };
+    }
+
+    pub const Function = fn (?*anyopaque) callconv(.C) void;
+};
+
+pub fn pushCleanupHook(
+    this: *RareData,
+    globalThis: *JSC.JSGlobalObject,
+    ctx: ?*anyopaque,
+    func: CleanupHook.Function,
+) void {
+    var hook = JSC.VirtualMachine.vm.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;
+    }
+
+    return hook;
+}
+
 pub fn boringEngine(rare: *RareData) *BoringSSL.ENGINE {
     return rare.boring_ssl_engine orelse brk: {
         rare.boring_ssl_engine = BoringSSL.ENGINE_new();
diff --git a/src/napi/napi.zig b/src/napi/napi.zig
index 74973806a..5a8ddc966 100644
--- a/src/napi/napi.zig
+++ b/src/napi/napi.zig
@@ -1077,10 +1077,48 @@ pub export fn napi_get_node_version(_: napi_env, version: **const napi_node_vers
     version.* = &napi_node_version.global;
     return .ok;
 }
-pub extern fn napi_get_uv_event_loop(env: napi_env, loop: [*]*struct_uv_loop_s) napi_status;
+pub export fn napi_get_uv_event_loop(_: napi_env, loop: *?*struct_uv_loop_s) napi_status {
+    // lol
+    loop.* = JSC.VirtualMachine.vm.eventLoop();
+}
 pub extern fn napi_fatal_exception(env: napi_env, err: napi_value) napi_status;
-pub extern fn napi_add_env_cleanup_hook(env: napi_env, fun: ?fn (?*anyopaque) callconv(.C) void, arg: ?*anyopaque) napi_status;
-pub extern fn napi_remove_env_cleanup_hook(env: napi_env, fun: ?fn (?*anyopaque) callconv(.C) void, arg: ?*anyopaque) napi_status;
+
+// We use a linked list here because we assume removing these is relatively rare
+// and array reallocations are relatively expensive.
+pub export fn napi_add_env_cleanup_hook(env: napi_env, fun: ?fn (?*anyopaque) callconv(.C) void, arg: ?*anyopaque) napi_status {
+    if (fun == null)
+        return .ok;
+
+    JSC.VirtualMachine.vm.rareData().pushCleanupHook(env, arg, fun);
+    return .ok;
+}
+pub export fn napi_remove_env_cleanup_hook(env: napi_env, fun: ?fn (?*anyopaque) callconv(.C) void, arg: ?*anyopaque) napi_status {
+    if (JSC.VirtualMachine.vm.rare_data == null or fun == null)
+        return .ok;
+
+    var rare_data = JSC.VirtualMachine.vm.rare_data.?;
+    var hook = rare_data.cleanup_hook orelse return .ok;
+    const cmp = JSC.RareData.CleanupHook.from(env, arg, fun.?);
+    if (hook.eql(cmp)) {
+        JSC.VirtualMachine.vm.allocator.destroy(hook);
+        rare_data.cleanup_hook = null;
+        rare_data.tail_cleanup_hook = null;
+    }
+    while (hook.next) |current| {
+        if (hook.eql(cmp)) {
+            if (current.next) |next| {
+                hook.next = next;
+            } else {
+                hook.next = null;
+            }
+            JSC.VirtualMachine.vm.allocator.destroy(current);
+            return .ok;
+        }
+        hook = current;
+    }
+
+    return .ok;
+}
 pub extern fn napi_open_callback_scope(env: napi_env, resource_object: napi_value, context: napi_async_context, result: [*c]napi_callback_scope) napi_status;
 pub extern fn napi_close_callback_scope(env: napi_env, scope: napi_callback_scope) napi_status;
 pub extern fn napi_create_threadsafe_function(env: napi_env, func: napi_value, async_resource: napi_value, async_resource_name: napi_value, max_queue_size: usize, initial_thread_count: usize, thread_finalize_data: ?*anyopaque, thread_finalize_cb: napi_finalize, context: ?*anyopaque, call_js_cb: napi_threadsafe_function_call_js, result: [*c]napi_threadsafe_function) napi_status;
@@ -1090,8 +1128,15 @@ pub extern fn napi_acquire_threadsafe_function(func: napi_threadsafe_function) n
 pub extern fn napi_release_threadsafe_function(func: napi_threadsafe_function, mode: napi_threadsafe_function_release_mode) napi_status;
 pub extern fn napi_unref_threadsafe_function(env: napi_env, func: napi_threadsafe_function) napi_status;
 pub extern fn napi_ref_threadsafe_function(env: napi_env, func: napi_threadsafe_function) napi_status;
-pub extern fn napi_add_async_cleanup_hook(env: napi_env, hook: napi_async_cleanup_hook, arg: ?*anyopaque, remove_handle: [*c]napi_async_cleanup_hook_handle) napi_status;
-pub extern fn napi_remove_async_cleanup_hook(remove_handle: napi_async_cleanup_hook_handle) napi_status;
+
+pub export fn napi_add_async_cleanup_hook(_: napi_env, _: napi_async_cleanup_hook, _: ?*anyopaque, _: [*c]napi_async_cleanup_hook_handle) napi_status {
+    // TODO:
+    return .ok;
+}
+pub export fn napi_remove_async_cleanup_hook(_: napi_async_cleanup_hook_handle) napi_status {
+    // TODO:
+    return .ok;
+}
 
 pub const NAPI_VERSION_EXPERIMENTAL = @import("std").zig.c_translation.promoteIntLiteral(c_int, 2147483647, .decimal);
 pub const NAPI_VERSION = @as(c_int, 8);