Fixes #3512 (#3526)

* Fixes #3512

* Fix `clearTimeout` and `clearInterval` not cancelling jobs same-tick

---------

Co-authored-by: Jarred Sumner <709451+Jarred-Sumner@users.noreply.github.com>

9 files changed, 168 insertions, 38 deletions

diff --git a/packages/bun-types/timers.d.ts b/packages/bun-types/timers.d.ts
index ab1e29953..0d2f3e745 100644
--- a/packages/bun-types/timers.d.ts
+++ b/packages/bun-types/timers.d.ts
@@ -11,8 +11,8 @@
 
 declare module "timers" {
   class Timer {
-    ref(): void;
-    unref(): void;
+    ref(): Timer;
+    unref(): Timer;
     hasRef(): boolean;
   }
 
diff --git a/src/bun.js/api/bun.zig b/src/bun.js/api/bun.zig
index 1e5a5e004..fbf567446 100644
--- a/src/bun.js/api/bun.zig
+++ b/src/bun.js/api/bun.zig
@@ -3715,21 +3715,32 @@ pub const Timer = struct {
             const kind = this.kind;
             var map: *TimeoutMap = vm.timer.maps.get(kind);
 
-            // This doesn't deinit the timer
-            // Timers are deinit'd separately
-            // We do need to handle when the timer is cancelled after the job has been enqueued
-            if (kind != .setInterval) {
-                if (map.fetchSwapRemove(this.id) == null) {
-                    // if the timeout was cancelled, don't run the callback
-                    this.deinit();
-                    return;
-                }
-            } else {
-                if (!map.contains(this.id)) {
-                    // if the interval was cancelled, don't run the callback
-                    this.deinit();
-                    return;
+            const should_cancel_job = brk: {
+                // This doesn't deinit the timer
+                // Timers are deinit'd separately
+                // We do need to handle when the timer is cancelled after the job has been enqueued
+                if (kind != .setInterval) {
+                    if (map.get(this.id)) |tombstone_or_timer| {
+                        break :brk tombstone_or_timer != null;
+                    } else {
+                        // clearTimeout has been called
+                        break :brk true;
+                    }
+                } else {
+                    if (map.get(this.id)) |tombstone_or_timer| {
+                        // .refresh() was called after CallbackJob enqueued
+                        break :brk tombstone_or_timer == null;
+                    }
                 }
+
+                break :brk false;
+            };
+
+            if (should_cancel_job) {
+                this.deinit();
+                return;
+            } else if (kind != .setInterval) {
+                _ = map.swapRemove(this.id);
             }
 
             var args_buf: [8]JSC.JSValue = undefined;
@@ -3825,10 +3836,29 @@ pub const Timer = struct {
             return timer_js;
         }
 
-        pub fn doRef(this: *TimerObject, _: *JSC.JSGlobalObject, _: *JSC.CallFrame) callconv(.C) JSValue {
+        pub fn doRef(this: *TimerObject, globalObject: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) callconv(.C) JSValue {
+            const this_value = callframe.this();
+            this_value.ensureStillAlive();
             if (this.ref_count > 0)
                 this.ref_count +|= 1;
-            return JSValue.jsUndefined();
+
+            var vm = globalObject.bunVM();
+            switch (this.kind) {
+                .setTimeout, .setImmediate, .setInterval => {
+                    if (vm.timer.maps.get(this.kind).getPtr(this.id)) |val_| {
+                        if (val_.*) |*val| {
+                            val.poll_ref.ref(vm);
+
+                            if (val.did_unref_timer) {
+                                val.did_unref_timer = false;
+                                vm.uws_event_loop.?.num_polls += 1;
+                            }
+                        }
+                    }
+                },
+            }
+
+            return this_value;
         }
 
         pub fn doRefresh(this: *TimerObject, globalThis: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) callconv(.C) JSValue {
@@ -3924,20 +3954,27 @@ pub const Timer = struct {
             return JSValue.jsUndefined();
         }
 
-        pub fn doUnref(this: *TimerObject, globalObject: *JSC.JSGlobalObject, _: *JSC.CallFrame) callconv(.C) JSValue {
+        pub fn doUnref(this: *TimerObject, globalObject: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) callconv(.C) JSValue {
+            const this_value = callframe.this();
+            this_value.ensureStillAlive();
             this.ref_count -|= 1;
-            if (this.ref_count == 0) {
-                switch (this.kind) {
-                    .setTimeout, .setImmediate => {
-                        _ = clearTimeout(globalObject, JSValue.jsNumber(this.id));
-                    },
-                    .setInterval => {
-                        _ = clearInterval(globalObject, JSValue.jsNumber(this.id));
-                    },
-                }
+            var vm = globalObject.bunVM();
+            switch (this.kind) {
+                .setTimeout, .setImmediate, .setInterval => {
+                    if (vm.timer.maps.get(this.kind).getPtr(this.id)) |val_| {
+                        if (val_.*) |*val| {
+                            val.poll_ref.unref(vm);
+
+                            if (!val.did_unref_timer) {
+                                val.did_unref_timer = true;
+                                vm.uws_event_loop.?.num_polls -= 1;
+                            }
+                        }
+                    }
+                },
             }
 
-            return JSValue.jsUndefined();
+            return this_value;
         }
         pub fn hasRef(this: *TimerObject, globalObject: *JSC.JSGlobalObject, _: *JSC.CallFrame) callconv(.C) JSValue {
             return JSValue.jsBoolean(this.ref_count > 0 and globalObject.bunVM().timer.maps.get(this.kind).contains(this.id));
@@ -3959,6 +3996,7 @@ pub const Timer = struct {
         callback: JSC.Strong = .{},
         globalThis: *JSC.JSGlobalObject,
         timer: *uws.Timer,
+        did_unref_timer: bool = false,
         poll_ref: JSC.PollRef = JSC.PollRef.init(),
         arguments: JSC.Strong = .{},
 
@@ -4060,8 +4098,14 @@ pub const Timer = struct {
 
             var vm = this.globalThis.bunVM();
 
-            this.poll_ref.unrefOnNextTick(vm);
+            this.poll_ref.unref(vm);
+
             this.timer.deinit();
+            if (this.did_unref_timer) {
+                // balance double-unrefing
+                vm.uws_event_loop.?.num_polls += 1;
+            }
+
             this.callback.deinit();
             this.arguments.deinit();
         }
diff --git a/test/js/node/timers/node-timers.test.ts b/test/js/node/timers/node-timers.test.ts
index e6fa48010..412eabc22 100644
--- a/test/js/node/timers/node-timers.test.ts
+++ b/test/js/node/timers/node-timers.test.ts
@@ -1,17 +1,18 @@
 import { describe, test } from "bun:test";
-import { setTimeout, clearTimeout, setInterval, setImmediate } from "node:timers";
+import { setTimeout, clearTimeout, setInterval, clearInterval, setImmediate } from "node:timers";
 
-for (const fn of [setTimeout, setInterval, setImmediate]) {
+for (const fn of [setTimeout, setInterval]) {
   describe(fn.name, () => {
     test("unref is possible", done => {
       const timer = fn(() => {
         done(new Error("should not be called"));
-      }, 1);
-      fn(() => {
+      }, 1).unref();
+      const other = fn(() => {
+        clearInterval(other);
         done();
       }, 2);
-      timer.unref();
-      if (fn !== setImmediate) clearTimeout(timer);
+      if (fn === setTimeout) clearTimeout(timer);
+      if (fn === setInterval) clearInterval(timer);
     });
   });
 }
diff --git a/test/js/web/timers/setTimeout-unref-fixture-2.js b/test/js/web/timers/setTimeout-unref-fixture-2.js
new file mode 100644
index 000000000..6a78f13cd
--- /dev/null
+++ b/test/js/web/timers/setTimeout-unref-fixture-2.js
@@ -0,0 +1,9 @@
+setTimeout(() => {
+  console.log("TEST FAILED!");
+}, 100)
+  .ref()
+  .unref();
+
+setTimeout(() => {
+  // this one should always run
+}, 1);
diff --git a/test/js/web/timers/setTimeout-unref-fixture-3.js b/test/js/web/timers/setTimeout-unref-fixture-3.js
new file mode 100644
index 000000000..41808f5fc
--- /dev/null
+++ b/test/js/web/timers/setTimeout-unref-fixture-3.js
@@ -0,0 +1,7 @@
+setTimeout(() => {
+  setTimeout(() => {}, 999_999);
+}, 100).unref();
+
+setTimeout(() => {
+  // this one should always run
+}, 1);
diff --git a/test/js/web/timers/setTimeout-unref-fixture-4.js b/test/js/web/timers/setTimeout-unref-fixture-4.js
new file mode 100644
index 000000000..9968f3b36
--- /dev/null
+++ b/test/js/web/timers/setTimeout-unref-fixture-4.js
@@ -0,0 +1,5 @@
+setTimeout(() => {
+  console.log("TEST PASSED!");
+}, 1)
+  .unref()
+  .ref();
diff --git a/test/js/web/timers/setTimeout-unref-fixture-5.js b/test/js/web/timers/setTimeout-unref-fixture-5.js
new file mode 100644
index 000000000..e5caa1be4
--- /dev/null
+++ b/test/js/web/timers/setTimeout-unref-fixture-5.js
@@ -0,0 +1,5 @@
+setTimeout(() => {
+  console.log("TEST FAILED!");
+}, 100)
+  .ref()
+  .unref();
diff --git a/test/js/web/timers/setTimeout-unref-fixture.js b/test/js/web/timers/setTimeout-unref-fixture.js
new file mode 100644
index 000000000..97a0f78a2
--- /dev/null
+++ b/test/js/web/timers/setTimeout-unref-fixture.js
@@ -0,0 +1,12 @@
+const timer = setTimeout(() => {}, 999_999_999);
+if (timer.unref() !== timer) throw new Error("Expected timer.unref() === timer");
+
+var ranCount = 0;
+const going2Refresh = setTimeout(() => {
+  if (ranCount < 1) going2Refresh.refresh();
+  ranCount++;
+
+  if (ranCount === 2) {
+    console.log("SUCCESS");
+  }
+}, 1);
diff --git a/test/js/web/timers/setTimeout.test.js b/test/js/web/timers/setTimeout.test.js
index dbe89dea8..eef6bbae0 100644
--- a/test/js/web/timers/setTimeout.test.js
+++ b/test/js/web/timers/setTimeout.test.js
@@ -1,5 +1,7 @@
+import { spawnSync } from "bun";
 import { it, expect } from "bun:test";
-
+import { bunEnv, bunExe } from "harness";
+import path from "node:path";
 it("setTimeout", async () => {
   var lastID = -1;
   const result = await new Promise((resolve, reject) => {
@@ -172,11 +174,56 @@ it.skip("order of setTimeouts", done => {
   Promise.resolve().then(maybeDone(() => nums.push(1)));
 });
 
+it("setTimeout -> refresh", () => {
+  const { exitCode, stdout } = spawnSync({
+    cmd: [bunExe(), path.join(import.meta.dir, "setTimeout-unref-fixture.js")],
+    env: bunEnv,
+  });
+  expect(exitCode).toBe(0);
+  expect(stdout.toString()).toBe("SUCCESS\n");
+});
+
+it("setTimeout -> unref -> ref works", () => {
+  const { exitCode, stdout } = spawnSync({
+    cmd: [bunExe(), path.join(import.meta.dir, "setTimeout-unref-fixture-4.js")],
+    env: bunEnv,
+  });
+  expect(exitCode).toBe(0);
+  expect(stdout.toString()).toBe("TEST PASSED!\n");
+});
+
+it("setTimeout -> ref -> unref works, even if there is another timer", () => {
+  const { exitCode, stdout } = spawnSync({
+    cmd: [bunExe(), path.join(import.meta.dir, "setTimeout-unref-fixture-2.js")],
+    env: bunEnv,
+  });
+  expect(exitCode).toBe(0);
+  expect(stdout.toString()).toBe("");
+});
+
+it("setTimeout -> ref -> unref works", () => {
+  const { exitCode, stdout } = spawnSync({
+    cmd: [bunExe(), path.join(import.meta.dir, "setTimeout-unref-fixture-5.js")],
+    env: bunEnv,
+  });
+  expect(exitCode).toBe(0);
+  expect(stdout.toString()).toBe("");
+});
+
+it("setTimeout -> unref doesn't keep event loop alive forever", () => {
+  const { exitCode, stdout } = spawnSync({
+    cmd: [bunExe(), path.join(import.meta.dir, "setTimeout-unref-fixture-3.js")],
+    env: bunEnv,
+  });
+  expect(exitCode).toBe(0);
+  expect(stdout.toString()).toBe("");
+});
+
 it("setTimeout should refresh N times", done => {
   let count = 0;
   let timer = setTimeout(() => {
     count++;
-    timer.refresh();
+    expect(timer.refresh()).toBe(timer);
   }, 50);
 
   setTimeout(() => {