Merge branch 'main' into jarred/request-finalizer

src/bun.js/api/server.zig

@@ -921,6 +921,7 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
         method: HTTP.Method,
         aborted: bool = false,
         finalized: bun.DebugOnly(bool) = bun.DebugOnlyDefault(false),
+        request_object_finalization_callback: ?*JSC.FinalizationCallback = null,
         upgrade_context: ?*uws.uws_socket_context_t = null,
 
         /// We can only safely free once the request body promise is finalized
@@ -975,6 +976,11 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
             this.resp.onAborted(*RequestContext, RequestContext.onAbort, this);
         }
 
+        pub fn onJSRequestObjectFinalized(this: *RequestContext) void {
+            this.request_js_object = null;
+            this.request_object_finalization_callback = null;
+        }
+
         pub fn onResolve(_: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) callconv(.C) JSValue {
             ctxLog("onResolve", .{});
 
@@ -983,13 +989,13 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
             const result = arguments.ptr[0];
             result.ensureStillAlive();
 
+            // ctx.request_js_object is set to null by the finalization callback
             if (ctx.request_js_object != null and ctx.signal == null) {
                 var request_js = ctx.request_js_object.?.value();
                 request_js.ensureStillAlive();
                 if (request_js.as(Request)) |request_object| {
                     if (request_object.signal) |signal| {
-                        ctx.signal = signal;
-                        _ = signal.ref();
+                        ctx.signal = signal.ref();
                     }
                 }
             }
@@ -1042,13 +1048,13 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
             var ctx = arguments.ptr[1].asPromisePtr(@This());
             const err = arguments.ptr[0];
 
+            // ctx.request_js_object is set to null by the finalization callback
             if (ctx.request_js_object != null and ctx.signal == null) {
                 var request_js = ctx.request_js_object.?.value();
                 request_js.ensureStillAlive();
                 if (request_js.as(Request)) |request_object| {
                     if (request_object.signal) |signal| {
-                        ctx.signal = signal;
-                        _ = signal.ref();
+                        ctx.signal = signal.ref();
                     }
                 }
             }
@@ -1327,6 +1333,11 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
                     // User called .blob(), .json(), text(), or .arrayBuffer() on the Request object
                     // but we received nothing or the connection was aborted
                     if (request_js.as(Request)) |req| {
+                        if (this.request_object_finalization_callback) |finalizer| {
+                            finalizer.detach(this.server.vm.finalizationPool());
+                            req.finalization_callback = null;
+                            this.request_object_finalization_callback = null;
+                        }
 
                         // the promise is pending
                         if (req.body == .Locked and (req.body.Locked.action != .none or req.body.Locked.promise != null)) {
@@ -1411,11 +1422,21 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
                 // User called .blob(), .json(), text(), or .arrayBuffer() on the Request object
                 // but we received nothing or the connection was aborted
                 if (request_js.as(Request)) |req| {
+                    if (this.request_object_finalization_callback) |finalizer| {
+                        finalizer.detach(this.server.vm.finalizationPool());
+                        if (finalizer == this.request_object_finalization_callback)
+                            req.finalization_callback = null;
+
+                        this.request_object_finalization_callback = null;
+                    }
+
                     // the promise is pending
                     if (req.body == .Locked and req.body.Locked.action != .none and req.body.Locked.promise != null) {
                         req.body.toErrorInstance(JSC.toTypeError(.ABORT_ERR, "Request aborted", .{}, this.server.globalThis), this.server.globalThis);
                     }
+
                     req.uws_request = null;
+                    req.upgrader = null;
                 }
             }
 
@@ -2118,6 +2139,19 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
                 // so we have to clear it in here too
                 request_object.uws_request = null;
 
+                if (ctx.signal == null) {
+                    if (request_object.signal) |signal| {
+                        ctx.signal = signal.ref();
+                    } else if (ctx.request_object_finalization_callback == null and request_object.finalization_callback == null) {
+                        request_object.finalization_callback = JSC.FinalizationCallback.create(
+                            this.vm.finalizationPool(),
+                            RequestContext,
+                            ctx,
+                            RequestContext.onJSRequestObjectFinalized,
+                        );
+                    }
+                }
+
                 ctx.setAbortHandler();
                 ctx.pending_promises_for_abort += 1;
 

src/bun.js/base.zig

@@ -4144,3 +4144,94 @@ pub const BinaryType = enum {
         }
     }
 };
+
+pub const FinalizationCallback = struct {
+    context: Ptr = Ptr{},
+    callback: ?*const fn (*anyopaque) void = null,
+
+    pub const Ptr = packed struct(u64) {
+        address: u60 = 0,
+
+        // Really should only have 2 references maximum
+        ref_count: u4 = 0,
+
+        pub fn ref(this: *Ptr) void {
+            this.ref_count += 1;
+        }
+
+        pub fn unref(this: *Ptr) void {
+            this.ref_count -= 1;
+        }
+
+        pub fn ptr(this: Ptr) ?*anyopaque {
+            if (this.address == 0)
+                return null;
+
+            // zero the trailing little endian 4 bits from u64
+            // We want to make sure that the ref_count is not included in the pointer
+            const safe_addr = @as(usize, this.address) & @as(usize, 0xFFFFFFFFFFFFFFF0);
+            return @intToPtr(*anyopaque, safe_addr);
+        }
+
+        pub fn init(addr: *anyopaque) Ptr {
+            return .{
+                .address = @truncate(u60, @ptrToInt(addr)),
+                .ref_count = 1,
+            };
+        }
+    };
+
+    pub const Pool = bun.HiveArray(FinalizationCallback, 512).Fallback;
+
+    pub fn call(this: *FinalizationCallback, pool: *Pool) void {
+        var callback = this.callback orelse {
+            this.unregister();
+            pool.put(this);
+            return;
+        };
+        var ctx = this.context.ptr() orelse {
+            this.unregister();
+            pool.put(this);
+            return;
+        };
+        this.unregister();
+        pool.put(this);
+
+        callback(ctx);
+    }
+
+    /// Pointers to FinalizationCallback are invalidated immediately **before** the callback is called.
+    pub fn create(pool: *Pool, comptime Context: type, context: *Context, callback: *const fn (*Context) void) *FinalizationCallback {
+        var this: *FinalizationCallback = pool.get();
+        this.set(Context, context, callback);
+        return this;
+    }
+
+    pub fn set(this: *FinalizationCallback, comptime Context: type, context: *Context, callback: *const fn (*Context) void) void {
+        this.* = .{
+            .context = Ptr.init(bun.cast(*anyopaque, context)),
+            .callback = bun.cast(*const fn (*anyopaque) void, callback),
+        };
+    }
+
+    fn unregister(this: *FinalizationCallback) void {
+        this.* = .{};
+    }
+
+    pub fn ref(this: *FinalizationCallback) void {
+        this.context.ref();
+    }
+
+    pub fn detach(this: *FinalizationCallback, pool: *Pool) void {
+        this.unregister();
+        this.unref(pool);
+    }
+
+    pub fn unref(this: *FinalizationCallback, pool: *Pool) void {
+        this.context.unref();
+        if (this.context.ref_count == 0) {
+            this.unregister();
+            pool.put(this);
+        }
+    }
+};

src/bun.js/javascript.zig

@@ -361,6 +361,7 @@ pub const VirtualMachine = struct {
     pending_unref_counter: i32 = 0,
     preload: []const string = &[_][]const u8{},
     unhandled_pending_rejection_to_capture: ?*JSC.JSValue = null,
+    finalization_pool: ?*JSC.FinalizationCallback.Pool = null,
 
     hot_reload: bun.CLI.Command.HotReload = .none,
 
@@ -464,6 +465,16 @@ pub const VirtualMachine = struct {
 
     pub threadlocal var is_main_thread_vm: bool = false;
 
+    pub fn finalizationPool(this: *VirtualMachine) *JSC.FinalizationCallback.Pool {
+        if (this.finalization_pool == null) {
+            var pool = this.allocator.create(JSC.FinalizationCallback.Pool) catch @panic("Failed to create finalization pool");
+            pool.* = JSC.FinalizationCallback.Pool.init(this.allocator);
+            this.finalization_pool = pool;
+        }
+
+        return this.finalization_pool.?;
+    }
+
     pub const UnhandledRejectionScope = struct {
         ctx: ?*anyopaque = null,
         onUnhandledRejection: *const OnUnhandledRejection = undefined,

src/bun.js/webcore/request.zig

@@ -68,6 +68,8 @@ pub const Request = struct {
     // We must report a consistent value for this
     reported_estimated_size: ?u63 = null,
 
+    finalization_callback: ?*JSC.FinalizationCallback = null,
+
     const RequestMixin = BodyMixin(@This());
     pub usingnamespace JSC.Codegen.JSRequest;
 
@@ -279,6 +281,10 @@ pub const Request = struct {
 
     pub fn finalize(this: *Request) callconv(.C) void {
         this.finalizeWithoutDeinit();
+        if (this.finalization_callback) |finalizer| {
+            var pool = JSC.VirtualMachine.get().finalizationPool();
+            finalizer.call(pool);
+        }
         bun.default_allocator.destroy(this);
     }
 

test/js/bun/http/serve.test.ts

@@ -123,11 +123,16 @@ it("request.signal works in trivial case", async () => {
 it("request.signal works in leaky case", async () => {
   var aborty = new AbortController();
   var didAbort = false;
-  var leaky: Request | undefined;
+  var onRequest = (req: Request) => {
+    req.signal.addEventListener("abort", () => {
+      didAbort = true;
+    });
+  };
+
   await runTest(
     {
       async fetch(req) {
-        leaky = req;
+        onRequest(req);
         expect(didAbort).toBe(false);
         aborty.abort();
         await Bun.sleep(2);
@@ -135,21 +140,9 @@ it("request.signal works in leaky case", async () => {
       },
     },
     async server => {
-      try {
-        const resp = fetch(`http://${server.hostname}:${server.port}`, { signal: aborty.signal });
-
-        await Bun.sleep(1);
-
-        leaky!.signal.addEventListener("abort", () => {
-          didAbort = true;
-        });
-
-        await resp;
-
-        throw new Error("Expected fetch to throw");
-      } catch (e: any) {
-        expect(e.name).toBe("AbortError");
-      }
+      expect(async () => {
+        await fetch(`http://${server.hostname}:${server.port}`, { signal: aborty.signal });
+      }).toThrow("The operation was aborted.");
 
       await Bun.sleep(1);