fix unrefOnNextTick

8 files changed, 41 insertions, 21 deletions

diff --git a/src/bun.js/api/bun.zig b/src/bun.js/api/bun.zig
index bcea1005c..d8506b6b5 100644
--- a/src/bun.js/api/bun.zig
+++ b/src/bun.js/api/bun.zig
@@ -3029,7 +3029,7 @@ pub const Timer = struct {
 
             var vm = this.globalThis.bunVM();
 
-            this.poll_ref.unref(vm);
+            this.poll_ref.unrefOnNextTick(vm);
             this.timer.deinit();
             this.callback.deinit();
             this.arguments.deinit();
diff --git a/src/bun.js/api/bun/dns_resolver.zig b/src/bun.js/api/bun/dns_resolver.zig
index baf7b5b3d..92577d556 100644
--- a/src/bun.js/api/bun/dns_resolver.zig
+++ b/src/bun.js/api/bun/dns_resolver.zig
@@ -614,10 +614,6 @@ pub const GetAddrInfo = struct {
     };
 };
 
-
-
-
-
 pub const ResolveSrvInfoRequest = struct {
     const log = Output.scoped(.ResolveSrvInfoRequest, false);
 
@@ -631,11 +627,10 @@ pub const ResolveSrvInfoRequest = struct {
     pub fn init(
         cache: DNSResolver.SrvCacheHit,
         resolver: ?*DNSResolver,
-        name: [] const u8,
+        name: []const u8,
         globalThis: *JSC.JSGlobalObject,
         comptime cache_field: []const u8,
     ) !*ResolveSrvInfoRequest {
-
         var request = try globalThis.allocator().create(ResolveSrvInfoRequest);
         var hasher = std.hash.Wyhash.init(0);
         hasher.update(name);
@@ -700,7 +695,6 @@ pub const ResolveSrvInfoRequest = struct {
 
     pub fn onCaresComplete(this: *ResolveSrvInfoRequest, err_: ?c_ares.Error, timeout: i32, result: ?*c_ares.struct_ares_srv_reply) void {
         if (this.resolver_for_caching) |resolver| {
-
             if (this.cache.pending_cache) {
                 resolver.drainPendingSrvCares(
                     this.cache.pos_in_pending,
@@ -998,13 +992,10 @@ pub const SrvLookup = struct {
         return;
     }
 
-
-
     pub fn onComplete(this: *SrvLookup, result: JSC.JSValue) void {
         var promise = this.promise;
         var globalThis = this.globalThis;
         this.promise = .{};
-
         promise.resolve(globalThis, result);
         this.deinit();
     }
@@ -1169,7 +1160,7 @@ pub const DNSResolver = struct {
 
         return entry;
     }
-     fn getSrvKey(this: *DNSResolver, index: u8, comptime cache_name: []const u8) ResolveSrvInfoRequest.PendingCacheKey {
+    fn getSrvKey(this: *DNSResolver, index: u8, comptime cache_name: []const u8) ResolveSrvInfoRequest.PendingCacheKey {
         var cache: *SrvPendingCache = &@field(this, cache_name);
         std.debug.assert(!cache.available.isSet(index));
         const entry = cache.buffer[index];
@@ -1207,7 +1198,6 @@ pub const DNSResolver = struct {
 
         array.ensureStillAlive();
 
-
         while (pending) |value| {
             var new_global = value.globalThis;
             if (prev_global != new_global) {
@@ -1631,7 +1621,6 @@ pub const DNSResolver = struct {
             return .zero;
         }
 
-      
         const name = name_str.toSlice(globalThis, bun.default_allocator);
         defer name.deinit();
         var vm = globalThis.bunVM();
diff --git a/src/bun.js/api/server.zig b/src/bun.js/api/server.zig
index 4f97c7384..3612e7d3e 100644
--- a/src/bun.js/api/server.zig
+++ b/src/bun.js/api/server.zig
@@ -4528,7 +4528,7 @@ pub fn NewServer(comptime ssl_enabled_: bool, comptime debug_mode_: bool) type {
         pub fn unref(this: *ThisServer) void {
             if (!this.poll_ref.isActive()) return;
 
-            this.poll_ref.unref(this.vm);
+            this.poll_ref.unrefOnNextTick(this.vm);
             this.vm.eventLoop().start_server_on_next_tick = false;
         }
 
diff --git a/src/bun.js/base.zig b/src/bun.js/base.zig
index 7d36662b6..3499fc6fe 100644
--- a/src/bun.js/base.zig
+++ b/src/bun.js/base.zig
@@ -3232,7 +3232,7 @@ pub const PollRef = struct {
         if (this.status != .active)
             return;
         this.status = .inactive;
-        vm.uws_event_loop.?.nextTick(*uws.Loop, vm.uws_event_loop.?, uws.Loop.unref);
+        vm.pending_unref_counter +|= 1;
     }
 
     /// Allow a poll to keep the process alive.
diff --git a/src/bun.js/event_loop.zig b/src/bun.js/event_loop.zig
index 9529d6d39..8dbf5abfd 100644
--- a/src/bun.js/event_loop.zig
+++ b/src/bun.js/event_loop.zig
@@ -476,7 +476,21 @@ pub const EventLoop = struct {
     }
 
     pub fn autoTick(this: *EventLoop) void {
-        var loop = this.virtual_machine.uws_event_loop.?;
+        var ctx = this.virtual_machine;
+        var loop = ctx.uws_event_loop.?;
+
+        // Some tasks need to keep the event loop alive for one more tick.
+        // We want to keep the event loop alive long enough to process those ticks and any microtasks
+        //
+        // BUT. We don't actually have an idle event in that case.
+        // That means the process will be waiting forever on nothing.
+        // So we need to drain the counter immediately before entering uSockets loop
+        const pending_unref = ctx.pending_unref_counter;
+        if (pending_unref > 0) {
+            ctx.pending_unref_counter = 0;
+            loop.unrefCount(pending_unref);
+        }
+
         if (loop.num_polls > 0 or loop.active > 0) {
             loop.tick();
             this.processGCTimer();
@@ -486,6 +500,15 @@ pub const EventLoop = struct {
 
     pub fn autoTickActive(this: *EventLoop) void {
         var loop = this.virtual_machine.uws_event_loop.?;
+
+        var ctx = this.virtual_machine;
+
+        const pending_unref = ctx.pending_unref_counter;
+        if (pending_unref > 0) {
+            ctx.pending_unref_counter = 0;
+            loop.unrefCount(pending_unref);
+        }
+
         if (loop.active > 0) {
             loop.tick();
             this.processGCTimer();
@@ -497,7 +520,6 @@ pub const EventLoop = struct {
         this.virtual_machine.gc_controller.processGCTimer();
     }
 
-    // TODO: fix this technical debt
     pub fn tick(this: *EventLoop) void {
         var ctx = this.virtual_machine;
         this.tickConcurrent();
@@ -517,6 +539,8 @@ pub const EventLoop = struct {
             break;
         }
 
+        // TODO: unify the event loops
+        // This needs a hook into JSC to schedule timers
         this.global.vm().doWork();
 
         while (this.tickWithCount() > 0) {
diff --git a/src/bun.js/javascript.zig b/src/bun.js/javascript.zig
index 49f9df802..ba8a9b276 100644
--- a/src/bun.js/javascript.zig
+++ b/src/bun.js/javascript.zig
@@ -359,6 +359,7 @@ pub const VirtualMachine = struct {
     has_loaded_node_modules: bool = false,
     timer: Bun.Timer = Bun.Timer{},
     uws_event_loop: ?*uws.Loop = null,
+    pending_unref_counter: i32 = 0,
 
     /// hide bun:wrap from stack traces
     /// bun:wrap is very noisy
diff --git a/src/deps/uws.zig b/src/deps/uws.zig
index 177fc1973..8eb2aaba9 100644
--- a/src/deps/uws.zig
+++ b/src/deps/uws.zig
@@ -464,6 +464,12 @@ pub const Loop = extern struct {
         this.active -= 1;
     }
 
+    pub fn unrefCount(this: *Loop, count: i32) void {
+        log("unref x {d}", .{count});
+        this.num_polls -|= count;
+        this.active -|= @intCast(u32, count);
+    }
+
     pub fn get() ?*Loop {
         return uws_get_loop();
     }
diff --git a/src/http/websocket_http_client.zig b/src/http/websocket_http_client.zig
index 77c64cf46..6692b8c87 100644
--- a/src/http/websocket_http_client.zig
+++ b/src/http/websocket_http_client.zig
@@ -268,7 +268,7 @@ pub fn NewHTTPUpgradeClient(comptime ssl: bool) type {
             this.input_body_buf.len = 0;
         }
         pub fn clearData(this: *HTTPClient) void {
-            this.poll_ref.unref(JSC.VirtualMachine.get());
+            this.poll_ref.unrefOnNextTick(JSC.VirtualMachine.get());
 
             this.clearInput();
             if (this.body_buf) |buf| {
@@ -866,7 +866,7 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
         }
 
         pub fn clearData(this: *WebSocket) void {
-            this.poll_ref.unref(this.globalThis.bunVM());
+            this.poll_ref.unrefOnNextTick(this.globalThis.bunVM());
             this.clearReceiveBuffers(true);
             this.clearSendBuffers(true);
             this.ping_len = 0;
@@ -1453,7 +1453,7 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
 
         fn dispatchClose(this: *WebSocket) void {
             var out = this.outgoing_websocket orelse return;
-            this.poll_ref.unref(this.globalThis.bunVM());
+            this.poll_ref.unrefOnNextTick(this.globalThis.bunVM());
             JSC.markBinding(@src());
             WebSocket__didCloseWithErrorCode(out, ErrorCode.closed);
         }