bake: re-arm the memory visualizer timer when it fires

src/runtime/bake/DevServer.rs

@@ -1090,10 +1090,7 @@ impl Drop for DevServer {
             debug_assert!(self.active_websocket_connections.is_empty());
         }
 
-        if self.memory_visualizer_timer.state == EventLoopTimerState::ACTIVE {
-            let timer_ptr: *mut EventLoopTimer = &raw mut self.memory_visualizer_timer;
-            self.timer_heap().remove(timer_ptr);
-        }
+        self.disarm_memory_visualizer_timer();
         self.graph_safety_lock.lock();
         // Hand ownership of the heap allocation to the watcher thread (which frees it in
         // `thread_main` once `running` flips false). Auto-dropping the `Box`
@@ -5451,13 +5448,42 @@ impl DevServer {
         crate::jsc_hooks::timer_all_mut()
     }
 
-    pub fn emit_memory_visualizer_message_timer(
-        _timer: &mut EventLoopTimer,
-        _: &bun_core::Timespec,
-    ) {
+    const MEMORY_VISUALIZER_TICK_MS: i64 = 1000;
+
+    pub(crate) fn arm_memory_visualizer_timer(&mut self) {
+        let next = bun_core::Timespec::ms_from_now(
+            bun_core::TimespecMockMode::ForceRealTime,
+            Self::MEMORY_VISUALIZER_TICK_MS,
+        );
+        let timer_ptr: *mut EventLoopTimer = &raw mut self.memory_visualizer_timer;
+        self.timer_heap().update(timer_ptr, &next);
+    }
+
+    pub(crate) fn disarm_memory_visualizer_timer(&mut self) {
+        if self.memory_visualizer_timer.state != EventLoopTimerState::ACTIVE {
+            return;
+        }
+        let timer_ptr: *mut EventLoopTimer = &raw mut self.memory_visualizer_timer;
+        self.timer_heap().remove(timer_ptr);
+    }
+
+    /// SAFETY: `timer` must be the `memory_visualizer_timer` field of a live, boxed `DevServer`.
+    pub(crate) unsafe fn emit_memory_visualizer_message_timer(timer: *mut EventLoopTimer) {
+        // SAFETY: caller contract; `from_timer_ptr` recovers the owning DevServer.
+        let dev: &mut DevServer = unsafe { &mut *DevServer::from_timer_ptr(timer) };
+        debug_assert!(dev.magic == Magic::Valid);
+        // Already popped by the event loop; left ACTIVE, `disarm` would remove() it again.
+        dev.memory_visualizer_timer.state = EventLoopTimerState::FIRED;
+        dev.emit_memory_visualizer_message();
+        dev.arm_memory_visualizer_timer();
     }
 
-    pub fn emit_memory_visualizer_message_if_needed(&mut self) {}
+    pub fn emit_memory_visualizer_message_if_needed(&mut self) {
+        if self.emit_memory_visualizer_events == 0 {
+            return;
+        }
+        self.emit_memory_visualizer_message();
+    }
 
     pub fn emit_memory_visualizer_message(&mut self) {
         debug_assert!(self.emit_memory_visualizer_events > 0);

src/runtime/bake/dev_server/hmr_socket.rs

@@ -124,23 +124,7 @@ impl HmrSocket {
                                             dev.memory_visualizer_timer.state
                                                 != EventLoopTimerState::ACTIVE
                                         );
-                                        // Note (jsc/runtime crate cycle): `vm.timer` is `()` on the
-                                        // low-tier `VirtualMachine`; the real `timer::All`
-                                        // lives in `RuntimeState` (see jsc_hooks.rs).
-                                        let state = crate::jsc_hooks::runtime_state();
-                                        let next = bun_core::Timespec::ms_from_now(
-                                            bun_core::TimespecMockMode::ForceRealTime,
-                                            1000,
-                                        );
-                                        // SAFETY: `runtime_state()` is non-null after
-                                        // `bun_runtime::init()`; JS-thread only, sole
-                                        // `&mut` to `timer` in this scope.
-                                        unsafe {
-                                            (*state).timer.update(
-                                                &raw mut dev.memory_visualizer_timer,
-                                                &next,
-                                            );
-                                        }
+                                        dev.arm_memory_visualizer_timer();
                                     }
                                 }
                                 _ => {}
@@ -309,17 +293,8 @@ impl HmrSocket {
             }
             if field.contains(HmrTopic::MemoryVisualizer.as_bit()) {
                 dev.emit_memory_visualizer_events -= 1;
-                if dev.emit_incremental_visualizer_events == 0
-                    && dev.memory_visualizer_timer.state == EventLoopTimerState::ACTIVE
-                {
-                    // Note (jsc/runtime crate cycle): `vm.timer` is `()` on the low-tier
-                    // `VirtualMachine`; the real `timer::All` lives in `RuntimeState`.
-                    let state = crate::jsc_hooks::runtime_state();
-                    // SAFETY: `runtime_state()` is non-null after `bun_runtime::init()`;
-                    // JS-thread only, sole `&mut` to `timer` in this scope.
-                    unsafe {
-                        (*state).timer.remove(&raw mut dev.memory_visualizer_timer);
-                    }
+                if dev.emit_memory_visualizer_events == 0 {
+                    dev.disarm_memory_visualizer_timer();
                 }
             }
         }

src/runtime/dispatch.rs

@@ -1097,7 +1097,7 @@ pub(crate) unsafe fn __bun_fire_timer(
         EventLoopTimerTag::DevServerMemoryVisualizerTick => {
             // SAFETY: per fn contract; `t` is the `memory_visualizer_timer`
             // field of a live DevServer.
-            DevServer::emit_memory_visualizer_message_timer(unsafe { &mut *t }, unsafe { &*now });
+            unsafe { DevServer::emit_memory_visualizer_message_timer(t) };
             Ok(())
         }
         EventLoopTimerTag::BunTest => {

test/bake/dev/hot.test.ts

@@ -1,7 +1,8 @@
 // Hot tests ensure that the `import.meta.hot` interface is functional
 import { expect } from "bun:test";
+import { isDebug } from "harness";
 import { renameSync, unlinkSync, writeFileSync } from "node:fs";
-import { devTest, emptyHtmlFile } from "../bake-harness";
+import { Dev, devTest, emptyHtmlFile, WAIT_MULTIPLIER } from "../bake-harness";
 
 devTest("import.meta.hot.accept basic", {
   files: {
@@ -612,6 +613,166 @@ devTest("hot update frames are not delivered to application websocket topics", {
   },
 });
 
+// The `M` (memory visualizer) HMR topic only exists in builds with
+// `BAKE_DEBUGGING_FEATURES` (canary or debug). A stable release accepts the
+// subscription but never emits a frame, so there is nothing to observe there.
+const hasBakeDebuggingFeatures = isDebug || Bun.version_with_sha.includes("-canary.");
+
+const memoryVisualizerApp = {
+  "index.html": emptyHtmlFile({}),
+  "bun.app.ts": `
+    import html from "./index.html";
+    // Always armed in the dev server's timer heap; /tick answers on its next fire.
+    const waiters = [];
+    setInterval(() => {
+      for (const resolve of waiters.splice(0)) resolve();
+    }, 20);
+    export default {
+      static: {
+        "/": html,
+      },
+      async fetch(req) {
+        if (new URL(req.url).pathname === "/tick") {
+          await new Promise(resolve => waiters.push(resolve));
+          return new Response("ticked");
+        }
+        return new Response("Not Found", { status: 404 });
+      },
+    };
+  `,
+};
+
+/**
+ * Opens an extra `/_bun/hmr` socket and subscribes it to the memory
+ * visualizer topic (`s` = subscribe, `M` = topic). The server answers the
+ * subscription with one `M` frame immediately and then one per timer tick.
+ */
+async function subscribeMemoryVisualizer(dev: Dev) {
+  const ws = new WebSocket(dev.baseUrl + "/_bun/hmr");
+  ws.binaryType = "arraybuffer";
+  let open = false;
+  let frames = 0;
+  // Set when the socket errors or closes unexpectedly; fails the step being
+  // awaited at that moment and every step started afterwards.
+  let failure: Error | null = null;
+  // The step currently being awaited: socket progress resolves it, socket
+  // failure or the deadline rejects it.
+  let step: { what: string; done: () => boolean; resolve: () => void; reject: (err: Error) => void } | null = null;
+  const check = () => {
+    if (step === null || !step.done()) return;
+    const { resolve } = step;
+    step = null;
+    resolve();
+  };
+  const fail = (reason: string) => {
+    const err = new Error(
+      `${reason}${step ? ` while ${step.what}` : ""} (received ${frames} memory visualizer frames)`,
+    );
+    failure ??= err;
+    if (step === null) return;
+    const { reject } = step;
+    step = null;
+    reject(err);
+  };
+  ws.onopen = () => {
+    open = true;
+    check();
+  };
+  ws.onerror = () => fail("hmr socket errored");
+  ws.onclose = event => fail(`hmr socket closed unexpectedly with code ${event.code}`);
+  ws.onmessage = event => {
+    if (new Uint8Array(event.data as ArrayBuffer)[0] !== "M".charCodeAt(0)) return;
+    frames++;
+    check();
+  };
+  const waitUntil = (what: string, done: () => boolean) =>
+    new Promise<void>((resolve, reject) => {
+      if (failure) return reject(failure);
+      const deadline = setTimeout(() => fail("timed out"), 5_000 * WAIT_MULTIPLIER);
+      step = {
+        what,
+        done,
+        resolve: () => {
+          clearTimeout(deadline);
+          resolve();
+        },
+        reject: err => {
+          clearTimeout(deadline);
+          reject(err);
+        },
+      };
+      check();
+    });
+
+  const handle = {
+    get frames() {
+      return frames;
+    },
+    /** Resolves once this socket has received `count` memory visualizer frames in total. */
+    waitForFrames: (count: number) => waitUntil(`waiting for memory visualizer frame #${count}`, () => frames >= count),
+    /** Closes the socket, which unsubscribes it on the server, and waits for the close handshake. */
+    close: () =>
+      new Promise<void>((resolve, reject) => {
+        if (failure) return reject(failure);
+        ws.onerror = () => reject(new Error("hmr socket errored during the close handshake"));
+        ws.onclose = () => resolve();
+        ws.close();
+      }),
+  };
+
+  await waitUntil("opening the hmr socket", () => open);
+  ws.send("sM");
+  await handle.waitForFrames(1);
+  return handle;
+}
+
+if (hasBakeDebuggingFeatures) {
+  devTest("memory visualizer topic ticks while subscribed and unsubscribes without disturbing other timers", {
+    files: memoryVisualizerApp,
+    htmlFiles: [],
+    async test(dev) {
+      const start = performance.now();
+      const subscriber = await subscribeMemoryVisualizer(dev);
+      // Frame 1 is sent synchronously on subscribe and frame 2 by the timer
+      // armed at that point; frame 3 only arrives if the tick handler
+      // re-armed the timer after firing.
+      await subscriber.waitForFrames(3);
+      // Two real ticks take at least 2s. A handler that re-inserted the timer
+      // without moving its deadline would deliver frame 3 right after frame 2,
+      // i.e. at roughly 1s; 1.5s keeps clear of both.
+      expect(performance.now() - start).toBeGreaterThanOrEqual(1500);
+
+      // The fixture's interval is pending in the same heap while the subscriber
+      // leaves. Unsubscribing removes the visualizer timer from the heap; when
+      // the tick had left that node marked as still in the heap, the removal
+      // discarded whichever timer was at the root instead (debug builds assert),
+      // and the interval never fired again.
+      await dev.fetch("/tick").equals("ticked");
+      await subscriber.close();
+      await dev.fetch("/tick").equals("ticked");
+
+      // The timer can be re-armed after being disarmed. This socket is left
+      // open on purpose: the harness's graceful exit closes it while the timer
+      // is armed, covering the teardown path.
+      await subscribeMemoryVisualizer(dev);
+    },
+  });
+
+  devTest("memory visualizer timer stays armed while another subscriber remains", {
+    files: memoryVisualizerApp,
+    htmlFiles: [],
+    async test(dev) {
+      const first = await subscribeMemoryVisualizer(dev);
+      const second = await subscribeMemoryVisualizer(dev);
+      await first.close();
+      // A frame published by a tick just before the close can still be in
+      // flight, so only the second frame from here on proves the timer is
+      // still armed for the remaining subscriber.
+      await second.waitForFrames(second.frames + 2);
+    },
+  });
+}
+
 devTest("dev.write resolves only after the new module body has run", {
   files: {
     "index.html": emptyHtmlFile({ scripts: ["index.ts"] }),