timers: take the event loop ref in Immediate.ref() only while the Immediate is queued

src/event_loop/EventLoopTimer.rs

@@ -491,7 +491,7 @@ impl TimerFlags {
     }
     /// Set to `true` only during execution of the JavaScript function so that
     /// `_destroyed` can be false during the callback even though `state` will
-    /// be `FIRED`.
+    /// be `FIRED`. A `setImmediate` on the immediate queue never sets it.
     #[inline]
     pub fn in_callback(self) -> bool {
         self.0 & Self::IN_CALLBACK != 0

src/runtime/timer/timer_object_internals.rs

@@ -200,11 +200,16 @@ impl TimerObjectInternals {
         let delta = if enable { 1 } else { -1 };
         match self.flags.get().kind() {
             // setImmediate has slightly different event loop logic
-            // SAFETY: `state` points at the boxed per-thread `RuntimeState`;
-            // single-threaded JS heap so no concurrent `&mut` to `.timer`.
-            Kind::SetImmediate if !self.flags.get().fake() => unsafe {
-                (*state).timer.increment_immediate_ref(delta, uws_loop)
-            },
+            Kind::SetImmediate if !self.flags.get().fake() => {
+                // `init()` takes the ref before it pins the wrapper.
+                debug_assert!(
+                    !enable || self.is_queued_immediate() || self.this_value.get().is_empty(),
+                    "an Immediate that left the queue took an event loop ref"
+                );
+                // SAFETY: `state` points at the boxed per-thread `RuntimeState`;
+                // single-threaded JS heap so no concurrent `&mut` to `.timer`.
+                unsafe { (*state).timer.increment_immediate_ref(delta, uws_loop) }
+            }
             // SAFETY: as above.
             _ => unsafe { (*state).timer.increment_timer_ref(delta, uws_loop) },
         }
@@ -214,20 +219,9 @@ impl TimerObjectInternals {
     /// C++ `Bun__JSTimeout__call` thunk (which handles exceptions internally).
     /// Returns `true` if an exception was thrown.
     ///
-    /// Note (noalias re-entrancy): takes `&self`, NOT `&mut self`.
-    /// The JS callback can re-enter `cancel()`/`do_refresh()` on this same
-    /// object via another `&Self` derived from the JS wrapper's `m_ptr`.
-    /// With `&mut self` here, LLVM's `noalias` lets it keep `self.flags` in a
-    /// register across the FFI call, so `set_in_callback(false)`'s RMW
-    /// clobbers the `has_cleared_timer` bit that `cancel()` set — the interval
-    /// re-fires forever. The fields they write are `Cell`/`JsCell`, so `&self`
-    /// carries no such guarantee.
-    ///
     /// # Safety
-    /// `vm` is the live per-thread VM, `thrown` is null or a live slot, and the
-    /// caller's `ref_()` pins the parent container for the duration of the call.
+    /// `vm` is the live per-thread VM and `thrown` is null or a live slot.
     unsafe fn run(
-        &self,
         global_this: *mut JSGlobalObject,
         timer: JSValue,
         callback: JSValue,
@@ -245,16 +239,8 @@ impl TimerObjectInternals {
         }
 
         // Bun__JSTimeout__call handles exceptions.
-        // `Cell<Flags>` RMW so the `in_callback` write reaches memory before JS
-        // runs (re-entrant `_destroyed` getter reads it via a different pointer).
-        self.update_flags(|f| f.set_in_callback(true));
         // SAFETY: `thrown` is null or the caller's live stack slot.
         let result = unsafe { Bun__JSTimeout__call(global, timer, callback, arguments, thrown) };
-        // No early returns between the `in_callback` set and this clear.
-        // `Cell<Flags>` RMW: must reload `flags` from memory — re-entrant
-        // `cancel()` may have set `has_cleared_timer` / cleared
-        // `is_keeping_event_loop_alive`.
-        self.update_flags(|f| f.set_in_callback(false));
 
         // SAFETY: as above.
         if unsafe { (*vm).is_inspector_enabled() } {
@@ -371,7 +357,7 @@ impl TimerObjectInternals {
     /// `Self::run` re-enters JS which can `cancel()`/`do_refresh()` this same
     /// object via the JS wrapper's `m_ptr`. With `&mut self` LLVM may cache
     /// `self.flags`/`event_loop_timer().state` across the call and clobber the
-    /// re-entrant write (see `run()` doc). Use a raw
+    /// re-entrant write (see `fire_catching`). Use a raw
     /// pointer; helper calls `(*this).foo()` materialise short-lived `&mut`
     /// scoped to each statement only — none span the JS call.
     ///
@@ -438,7 +424,7 @@ impl TimerObjectInternals {
             // SAFETY: `vm` per fn contract; `ref_()` above pins the parent
             // across re-entrancy.
             let result = unsafe {
-                s.run(
+                Self::run(
                     global_this,
                     timer,
                     callback,
@@ -605,10 +591,12 @@ impl TimerObjectInternals {
             // block. Every path through the labelled-block + `is_timer_done`
             // tail reaches it (no `return` between here and the deref).
 
+            // No early return before the clear below, or `_destroyed` stays false.
+            s.update_flags(|f| f.set_in_callback(true));
             // SAFETY: `vm` and `thrown` per fn contract; `ref_()` above pins the
             // parent across re-entrancy.
             let _ = unsafe {
-                s.run(
+                Self::run(
                     global_this,
                     this_object,
                     callback,
@@ -618,6 +606,7 @@ impl TimerObjectInternals {
                     thrown,
                 )
             };
+            s.update_flags(|f| f.set_in_callback(false));
 
             // SAFETY: `state` is the boxed per-thread `RuntimeState`; field read only.
             let fake_clock_after = unsafe { (*state).timer.fake_timers.installs() };
@@ -990,19 +979,37 @@ impl TimerObjectInternals {
         self.update_flags(|f| f.set_has_js_ref(true));
 
         // https://github.com/nodejs/node/blob/a7cbb904745591c9a9d047a364c2c188e5470047/lib/internal/timers.js#L256
-        // and
-        // https://github.com/nodejs/node/blob/a7cbb904745591c9a9d047a364c2c188e5470047/lib/internal/timers.js#L685-L687
         // Node only re-enables the keep-alive ref when `!this._destroyed`. Checking
         // `has_cleared_timer` alone is not sufficient: a one-shot timer that has already fired
         // has `has_cleared_timer == false` but is still destroyed. Calling `.unref(); .ref()`
         // on such a timer would otherwise leak an event-loop ref and hang the process.
-        if !did_have_js_ref && !self.get_destroyed() {
-            self.set_enable_keeping_event_loop_alive(VirtualMachine::get_mut_ptr(), true);
+        if !did_have_js_ref {
+            self.keep_event_loop_alive_again();
         }
 
         this_value
     }
 
+    /// `ref()` after `unref()`. Out of line: `do_ref` stays small enough to inline.
+    #[inline(never)]
+    fn keep_event_loop_alive_again(&self) {
+        let flags = self.flags.get();
+        let still_runs = if flags.kind() == Kind::SetImmediate && !flags.fake() {
+            // Node keeps an idle unref'ed Immediate queued, so `ref()` runs it there. `run_immediate_task` already dropped it here.
+            self.is_queued_immediate()
+        } else {
+            !self.get_destroyed()
+        };
+        if still_runs {
+            self.set_enable_keeping_event_loop_alive(VirtualMachine::get_mut_ptr(), true);
+        }
+    }
+
+    /// Node's `kRefed !== null`. `init()` pins the wrapper of a `setImmediate`, and each way off the immediate queue unpins it.
+    fn is_queued_immediate(&self) -> bool {
+        self.this_value.get().is_strong()
+    }
+
     pub(crate) fn do_unref(&self, this_value: JSValue) -> JSValue {
         this_value.ensure_still_alive();
 

test/js/node/timers/node-timers.test.ts

@@ -1,5 +1,5 @@
 import jsc from "bun:jsc";
-import { describe, expect, it, mock, test } from "bun:test";
+import { describe, expect, it, jest, mock, test } from "bun:test";
 import { bunEnv, bunExe, bunRun, isWindows } from "harness";
 import { AsyncLocalStorage } from "node:async_hooks";
 import path from "node:path";
@@ -137,7 +137,7 @@ describe("_destroyed", () => {
   });
 
   it("is false during the callback", async () => {
-    for (const fn of [setTimeout, setInterval, setImmediate]) {
+    for (const fn of [setTimeout, setInterval]) {
       const { promise: done, resolve } = Promise.withResolvers();
       const timer = fn(() => {
         try {
@@ -152,6 +152,17 @@ describe("_destroyed", () => {
     }
   });
 
+  // Node.js destroys an Immediate before its callback and a Timeout after it.
+  it("is true during the callback of an immediate, where ref() does nothing", async () => {
+    const { promise, resolve } = Promise.withResolvers();
+    const immediate = setImmediate(() => {
+      const inCallback = { destroyed: immediate._destroyed, hasRef: immediate.hasRef() };
+      immediate.unref().ref();
+      resolve({ ...inCallback, hasRefAfterRef: immediate.hasRef() });
+    }) as TimerWithDestroyed;
+    expect(await promise).toEqual({ destroyed: true, hasRef: false, hasRefAfterRef: false });
+  });
+
   it("is true after clearing", () => {
     const timeout = setTimeout(() => {}, 0) as TimerWithDestroyed;
     clearTimeout(timeout);
@@ -213,6 +224,149 @@ describe("_destroyed", () => {
   });
 });
 
+// Each script prints `_destroyed` and `hasRef()`, then the process ends by itself.
+describe("ref() holds the event loop only for an Immediate that is still queued", () => {
+  async function run(script: string) {
+    await using proc = Bun.spawn({
+      cmd: [bunExe(), "-e", script],
+      env: bunEnv,
+      stdout: "pipe",
+      stderr: "pipe",
+      // bun:test does not kill the child of a concurrent test that times out.
+      timeout: 4_000,
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    return { stdout, stderr, signalCode: proc.signalCode, exitCode };
+  }
+
+  // Node.js prints the same.
+  it.concurrent.each([
+    [
+      "ref() inside its callback, after unref() before it",
+      `const immediate = setImmediate(() => {
+         console.log("callback", immediate._destroyed, immediate.hasRef());
+         immediate.ref();
+         console.log("ref", immediate._destroyed, immediate.hasRef());
+       });
+       setImmediate(() => console.log("later", immediate._destroyed, immediate.hasRef()));
+       immediate.unref();`,
+      "callback true false\nref true false\nlater true false\n",
+    ],
+    [
+      "unref() and ref() inside its callback",
+      `const immediate = setImmediate(() => {
+         immediate.unref();
+         console.log("unref", immediate._destroyed, immediate.hasRef());
+         immediate.ref();
+         console.log("ref", immediate._destroyed, immediate.hasRef());
+       });`,
+      "unref true false\nref true false\n",
+    ],
+    [
+      "this.ref() inside its callback",
+      `setImmediate(function () {
+         this.ref();
+         console.log("ref", this._destroyed, this.hasRef());
+       }).unref();
+       setImmediate(() => {});`,
+      "ref true false\n",
+    ],
+    [
+      "process.ref(immediate) inside its callback",
+      `const immediate = setImmediate(() => {
+         process.ref(immediate);
+         console.log("ref", immediate._destroyed, immediate.hasRef());
+       });
+       setImmediate(() => {});
+       immediate.unref();`,
+      "ref true false\n",
+    ],
+    [
+      "ref() inside its callback, which then throws",
+      `process.on("uncaughtException", error => {
+         console.log(error.message, immediate._destroyed, immediate.hasRef());
+       });
+       const immediate = setImmediate(() => {
+         immediate.ref();
+         throw new Error("thrown");
+       });
+       setImmediate(() => {});
+       immediate.unref();`,
+      "thrown true false\n",
+    ],
+    [
+      "ref() in the uncaughtException listener that its callback's throw calls",
+      `process.on("uncaughtException", error => {
+         console.log(error.message, immediate._destroyed, immediate.hasRef());
+         immediate.ref();
+         console.log("ref", immediate._destroyed, immediate.hasRef());
+       });
+       const immediate = setImmediate(() => {
+         throw new Error("thrown");
+       });
+       setImmediate(() => {});
+       immediate.unref();`,
+      "thrown true false\nref true false\n",
+    ],
+    [
+      "ref() on a queued Immediate, from the callback of another, runs it",
+      `setImmediate(() => {
+         queued.ref();
+         console.log("ref", queued._destroyed, queued.hasRef());
+       });
+       const queued = setImmediate(() => console.log("callback", queued._destroyed, queued.hasRef()));
+       queued.unref();`,
+      "ref false true\ncallback true false\n",
+    ],
+    [
+      "unref() then ref() before its callback runs it",
+      `const immediate = setImmediate(() => console.log("callback", immediate._destroyed, immediate.hasRef()));
+       immediate.unref().ref();
+       console.log("ref", immediate._destroyed, immediate.hasRef());`,
+      "ref false true\ncallback true false\n",
+    ],
+  ])("%s", async (_, script, stdout) => {
+    expect(await run(script)).toEqual({ stdout, stderr: "", signalCode: null, exitCode: 0 });
+  });
+
+  // Node.js keeps this Immediate queued: it prints "ref false true", then runs the callback.
+  it.concurrent("ref() in beforeExit does nothing on an unref'ed Immediate that the idle loop dropped", async () => {
+    const script = `
+      const immediate = setImmediate(() => console.log("callback", immediate._destroyed, immediate.hasRef()));
+      immediate.unref();
+      process.once("beforeExit", () => {
+        console.log("beforeExit", immediate._destroyed, immediate.hasRef());
+        immediate.ref();
+        console.log("ref", immediate._destroyed, immediate.hasRef());
+      });`;
+    expect(await run(script)).toEqual({
+      stdout: "beforeExit false false\nref false false\n",
+      stderr: "",
+      signalCode: null,
+      exitCode: 0,
+    });
+  });
+
+  // A faked Immediate is in the timer heap, which gives the ref back after the callback.
+  it("ref() inside the callback of a faked Immediate takes the ref", () => {
+    jest.useFakeTimers({ toFake: ["setImmediate"] });
+    try {
+      const hasRef: boolean[] = [];
+      const immediate = globalThis.setImmediate(() => {
+        immediate.unref();
+        hasRef.push(immediate.hasRef());
+        immediate.ref();
+        hasRef.push(immediate.hasRef());
+      });
+      jest.runAllTimers();
+      hasRef.push(immediate.hasRef());
+      expect({ hasRef, pending: jest.getTimerCount() }).toEqual({ hasRef: [false, true, false], pending: 0 });
+    } finally {
+      jest.useRealTimers();
+    }
+  });
+});
+
 describe("clear", () => {
   it("can clear the other kind of timer", async () => {
     const timeout1 = setTimeout(() => {