bun test: hold the event loop across a test callback's drain and inside fake timer controls

src/jsc/VirtualMachine.rs

@@ -1622,8 +1622,8 @@ impl VirtualMachine {
     }
 
     /// `event_loop().enter()` now, `.exit_without_checkpoint()` on drop, for a
-    /// dispatcher that drains microtasks itself: see
-    /// [`EventLoop::enter_scope_without_checkpoint`].
+    /// dispatcher that drains microtasks itself, or for native code that
+    /// script called: see [`EventLoop::enter_scope_without_checkpoint`].
     #[inline]
     pub fn enter_event_loop_scope_without_checkpoint(
         &self,

src/jsc/event_loop.rs

@@ -391,6 +391,10 @@ impl EventLoop {
     /// drains are unconditional, so the held count does not skip them, and the
     /// continuations they run are covered by it as well.
     ///
+    /// Also for native code that script called and that runs callbacks before
+    /// it returns (a fake timer control, `WebSocket::cancel`). It owes no
+    /// checkpoint: that one belongs to whoever called the script.
+    ///
     /// # Safety
     /// As [`Self::enter_scope`].
     #[inline]
@@ -1426,6 +1430,7 @@ impl EventLoop {
     /// Prefer `runCallbackWithResult` unless you really need to make sure that microtasks are drained.
     /// `context`: as for [`run_callback`](Self::run_callback).
     /// `after_call` runs after `callback` returns and before microtasks drain.
+    /// The loop stays entered until that drain is done: see [`Self::enter_scope_without_checkpoint`].
     pub fn run_callback_with_result_and_forcefully_drain_microtasks(
         &mut self,
         context: crate::ContextId,
@@ -1441,17 +1446,16 @@ impl EventLoop {
         };
         // R-2 noalias mitigation — see `run_callback` above.
         let this: *mut Self = core::hint::black_box(core::ptr::from_mut(self));
-        // SAFETY: `this` is the unique live `EventLoop`; short-lived `&mut`.
-        unsafe { (*this).enter() };
+        // SAFETY: `this` is the unique live `EventLoop`, which the VM owns for the process lifetime.
+        let _entered = unsafe { Self::enter_scope_without_checkpoint(this) };
         let result = callback.call(global_object, this_value, arguments);
-        let this: *mut Self = core::hint::black_box(this);
-        // SAFETY: as above.
-        unsafe { (*this).exit_without_checkpoint() };
         after_call();
         let result = result?;
         result.ensure_still_alive();
         let jsc_vm = global_object.bun_vm().jsc_vm();
-        self.drain_microtasks_with_global(global_object, jsc_vm)
+        let this: *mut Self = core::hint::black_box(this);
+        // SAFETY: as above; short-lived `&mut`.
+        unsafe { (*this).drain_microtasks_with_global(global_object, jsc_vm) }
             .map_err(|stopped| stopped.throw(global_object))?;
         Ok(result)
     }

src/runtime/test_runner/timers/FakeTimers.rs

@@ -1017,6 +1017,9 @@ impl FakeTimers {
         // SAFETY: per-thread `timer::All`; the borrow ends at this statement.
         unsafe { (*timer_all()).fake_timers.advance_clock_to(global, at) };
         let _firing = Firing::begin();
+        // A control fires this below the JS that called it, where the callback's exit must not drain microtasks.
+        // A drive step and the tick timer drain once they return.
+        let _entered = global.bun_vm().enter_event_loop_scope_without_checkpoint();
         let vm = VirtualMachine::get_mut_ptr();
         let mut thrown = JSValue::ZERO;
         // SAFETY: `timer` is live and its tag names its owner; the callees take raw pointers (noalias re-entrancy).
@@ -1052,6 +1055,12 @@ impl FakeTimers {
     fn run_ticks(global: &JSGlobalObject) -> Result<(), Stop> {
         let all = timer_all();
         // SAFETY: per-thread `timer::All`; each borrow lasts one statement and none spans a callback.
+        if unsafe { (*all).fake_timers.ticks.is_empty() } {
+            return Ok(());
+        }
+        // As in `fire`, for a native call that a tick makes and that runs a callback of its own.
+        let _entered = global.bun_vm().enter_event_loop_scope_without_checkpoint();
+        // SAFETY: as above.
         let limit = unsafe { (*all).fake_timers.loop_limit };
         let mut ran: u32 = 0;
         // SAFETY: as above.
@@ -1447,7 +1456,7 @@ impl FakeTimers {
         let TickMode::Interval(delta) = mode else {
             return Ok(());
         };
-        // Holds back the microtask drain until the whole advance is done, like sinon's `tick()`.
+        // Drains microtasks once the whole advance is done, like sinon's `tick()`.
         // SAFETY: the live event loop of the per-thread VM.
         let _entered =
             unsafe { bun_jsc::event_loop::EventLoop::enter_scope(global.bun_vm().event_loop()) };

src/runtime/test_runner/vi_wait.rs

@@ -215,6 +215,9 @@ impl ViWait {
         // SAFETY: heap-allocated by `to_js`; `finish` leaves the context before `root` releases the wrapper.
         unsafe { AbortHandle::arm_owner(this_ptr, cx.context()) };
 
+        // The first check runs below the JS that called this, so `waitFor()`, `waitUntil()` and `expect.poll()` count as
+        // timer controls: as in `FakeTimers::fire`, for a native call that `callback` makes.
+        let _entered = cx.vm().enter_event_loop_scope_without_checkpoint();
         if let Err(err) = this.poll(global, this_value, true) {
             match thrown_value(global, err) {
                 Ok(error) => this.settle(global, this_value, Err(error))?,

test/js/bun/test/bun_test.test.ts

@@ -1,4 +1,4 @@
-import { expect, test } from "bun:test";
+import { describe, expect, test } from "bun:test";
 import { bunEnv, bunExe, normalizeBunSnapshot, tempDir } from "harness";
 
 test("describe/test", async () => {
@@ -458,3 +458,175 @@ test("no --only flag with multiple files", async () => {
   `);
   expect(exitCode).toBe(0);
 });
+
+describe.concurrent("the callbacks of a test run with the event loop entered", () => {
+  async function printedBy(files: Record<string, string>) {
+    using dir = tempDir("bun-test-event-loop-entered", files);
+    await using proc = Bun.spawn({
+      cmd: [bunExe(), "test", `./${Object.keys(files)[0]}`],
+      env: bunEnv,
+      cwd: String(dir),
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    const printed = stdout
+      .split("\n")
+      .filter(line => line.startsWith("{"))
+      .map(line => JSON.parse(line));
+    return { printed, stderr, exitCode };
+  }
+
+  test("a native call in a promise job of a test returns before the next promise job runs", async () => {
+    const { printed, stderr, exitCode } = await printedBy({
+      "native-call.test.ts": `
+        import { afterAll, test } from "bun:test";
+        import { once } from "node:events";
+        import net from "node:net";
+
+        const opened = Promise.withResolvers();
+        const server = Bun.serve({
+          port: 0,
+          fetch: (request, server) => (server.upgrade(request, { data: {} }) ? undefined : new Response(null, { status: 400 })),
+          websocket: {
+            open: opened.resolve,
+            message() {},
+            close: socket => void socket.data.order.push("close handler"),
+          },
+        });
+        const client = new WebSocket("ws://localhost:" + server.port);
+        const serverSocket = await opened.promise;
+        const tcpServer = net.createServer(socket => socket.on("error", () => {})).listen(0, "127.0.0.1");
+        await once(tcpServer, "listening");
+        const tcpSocket = net.connect(tcpServer.address().port, "127.0.0.1");
+        await once(tcpSocket, "connect");
+        afterAll(() => {
+          server.stop(true);
+          tcpServer.close();
+        });
+
+        async function cell(where, call, native) {
+          const order = [];
+          await 0;
+          Promise.resolve().then(() => order.push("promise job"));
+          native(order);
+          order.push("returned");
+          await 0;
+          console.log(JSON.stringify({ where, call, order }));
+        }
+        // An HTMLRewriter handler runs in a promise job, and reader.cancel() settles a promise from native code.
+        async function cancelsReader(where) {
+          let controller;
+          const input = new ReadableStream({ start: c => void (controller = c) });
+          const order = [];
+          const response = new HTMLRewriter()
+            .on("p", {
+              element() {
+                order.push("handler start");
+                reader.cancel();
+                order.push("handler end");
+              },
+            })
+            .transform(new Response(input));
+          const reader = response.body.getReader();
+          const closed = reader.closed.then(() => order.push("reader.closed"));
+          controller.enqueue(new TextEncoder().encode("<p>x</p>"));
+          controller.close();
+          await closed;
+          console.log(JSON.stringify({ where, call: "reader.cancel()", order }));
+        }
+
+        // The file has loaded, and no test before these waits for the event loop.
+        const first = "the first tests of a file";
+        test(first + ": ServerWebSocket.close() calls the close handler", () =>
+          cell(first, "ServerWebSocket.close()", order => {
+            serverSocket.data.order = order;
+            serverSocket.close();
+          }));
+        test(first + ": net.Socket.destroy()", () => cell(first, "net.Socket.destroy()", () => tcpSocket.destroy()));
+        test(first + ": reader.cancel()", () => cancelsReader(first));
+        test("a test that waits for a real timer", () => new Promise(resolve => setTimeout(resolve, 1)));
+        test("a later test: reader.cancel()", () => cancelsReader("a later test"));
+      `,
+    });
+    const first = "the first tests of a file";
+    expect({ printed, exitCode }, stderr).toEqual({
+      printed: [
+        { where: first, call: "ServerWebSocket.close()", order: ["close handler", "returned", "promise job"] },
+        { where: first, call: "net.Socket.destroy()", order: ["returned", "promise job"] },
+        { where: first, call: "reader.cancel()", order: ["handler start", "handler end", "reader.closed"] },
+        { where: "a later test", call: "reader.cancel()", order: ["handler start", "handler end", "reader.closed"] },
+      ],
+      exitCode: 0,
+    });
+  });
+
+  // The task that reports a timeout runs the next steps, up to the first one that waits for the event loop.
+  // Their callbacks have one more entry below them.
+  test("in the steps that follow a test that timed out", async () => {
+    const { printed, stderr, exitCode } = await printedBy({
+      "after-timeout.test.ts": `
+        import { afterEach, describe, jest, test } from "bun:test";
+
+        test("times out", () => new Promise(() => {}), 1);
+
+        test("a fake timer control runs no promise job before it returns", async () => {
+          await 0;
+          jest.useFakeTimers();
+          const order = [];
+          setTimeout(async () => {
+            order.push("timer 1");
+            await 0;
+            order.push("after the await of timer 1");
+          }, 1);
+          setTimeout(() => order.push("timer 2"), 2);
+          jest.advanceTimersByTime(2);
+          order.push("returned");
+          jest.useRealTimers();
+          console.log(JSON.stringify({ order }));
+        });
+
+        describe("afterEach runs before a timer that is due", () => {
+          let where = "";
+          let fired = Promise.withResolvers();
+          const seen = [];
+          const callback = () => {
+            seen.push(where);
+            fired.resolve();
+          };
+          const untilTimersAreDue = () => {
+            const start = performance.now();
+            while (performance.now() - start < 2);
+          };
+          afterEach(async () => {
+            where = "afterEach";
+            await fired.promise;
+            fired = Promise.withResolvers();
+          });
+          test("when the test only waited for a promise job", async () => {
+            setTimeout(callback, 0);
+            untilTimersAreDue();
+            await 0;
+            where = "after the test";
+          });
+          test("when setImmediate resumed the test", async () => {
+            await new Promise(resolve => setImmediate(resolve));
+            setTimeout(callback, 0);
+            untilTimersAreDue();
+            where = "after the test";
+          });
+          test("prints", () => {
+            fired.resolve();
+            console.log(JSON.stringify({ seen }));
+          });
+        });
+      `,
+    });
+    expect({ printed, exitCode }, stderr).toEqual({
+      printed: [{ order: ["timer 1", "timer 2", "returned"] }, { seen: ["afterEach", "afterEach"] }],
+      exitCode: 1,
+    });
+    expect(stderr).toContain("(fail) times out");
+    expect(stderr).toContain(" 4 pass");
+  });
+});