bake: free the dev server's JS-thread bundle allocations with the bundle

src/bundler/bundle_v2.rs

@@ -6,6 +6,7 @@
 
 use core::ptr::NonNull;
 
+use bun_alloc::ast_alloc::{self, AstAllocState};
 use bun_collections::{ArrayHashMap, StringHashMap};
 use bun_core::ThreadLock;
 
@@ -140,6 +141,99 @@ pub struct BundleV2<'a> {
     /// dense and this is probed once per import in `on_parse_task_complete`
     /// (the main-thread parse-phase throughput limiter).
     pub(crate) requested_exports: Vec<Option<RequestedExports>>,
+
+    /// Declared last: `graph` / `linker` hold `AstVec`s backed by this state's
+    /// inline chunk, so it must drop after them.
+    pub(crate) async_ast_alloc: AsyncAstAlloc,
+}
+
+/// The `AstAllocState` of an `asynchronous` (dev server) bundle, spilling
+/// into `graph.heap`. The bundle's graph work runs as JS event loop callbacks,
+/// which would otherwise run with no state installed and leak every `AstAlloc`
+/// allocation on the global heap; each callback installs this one for its
+/// duration ([`BundleV2::enter_async_ast_scope`]).
+#[derive(Default)]
+pub(crate) struct AsyncAstAlloc(AsyncAstState);
+
+#[derive(Default)]
+enum AsyncAstState {
+    /// Synchronous bundle: the owning thread keeps its own state installed.
+    #[default]
+    Disabled,
+    Parked(Box<AstAllocState>),
+    Installed {
+        id: *const AstAllocState,
+        /// Thread-local occupant displaced by `enter`, reinstated by `exit`.
+        displaced: Option<Box<AstAllocState>>,
+    },
+}
+
+impl AsyncAstAlloc {
+    /// Returns the installed state's identity; null (and a no-op) unless parked.
+    fn enter(&mut self, spill: *mut bun_alloc::mimalloc::Heap) -> *const AstAllocState {
+        let mut state = match core::mem::take(&mut self.0) {
+            AsyncAstState::Parked(state) => state,
+            other => {
+                self.0 = other;
+                return core::ptr::null();
+            }
+        };
+        state.set_spill_heap(spill);
+        let id: *const AstAllocState = &raw const *state;
+        self.0 = AsyncAstState::Installed {
+            id,
+            displaced: ast_alloc::swap_state(Some(state)),
+        };
+        id
+    }
+
+    /// Uninstalls and parks the state again; no-op unless installed.
+    fn exit(&mut self) {
+        let (id, displaced) = match core::mem::take(&mut self.0) {
+            AsyncAstState::Installed { id, displaced } => (id, displaced),
+            other => {
+                self.0 = other;
+                return;
+            }
+        };
+        debug_assert!(
+            core::ptr::eq(ast_alloc::active_state_id(), id),
+            "AsyncAstAlloc::exit: another AstAllocState is still installed on top of the bundle's"
+        );
+        if let Some(state) = ast_alloc::swap_state(displaced) {
+            self.0 = AsyncAstState::Parked(state);
+        }
+    }
+}
+
+impl Drop for AsyncAstAlloc {
+    fn drop(&mut self) {
+        self.exit();
+    }
+}
+
+/// Uninstalls the bundle's state on drop. The callback holding it may complete
+/// the bundle, which frees the `BundleV2` before the callback returns; that
+/// teardown uninstalls the state itself, so `bv2` is only touched while the
+/// state this guard installed is still the active one.
+#[must_use = "the state is uninstalled as soon as the guard drops"]
+pub(crate) struct AsyncAstScope {
+    bv2: *mut BundleV2<'static>,
+    /// Null when this guard installed nothing.
+    installed: *const AstAllocState,
+}
+
+impl Drop for AsyncAstScope {
+    fn drop(&mut self) {
+        if self.installed.is_null() || !core::ptr::eq(ast_alloc::active_state_id(), self.installed)
+        {
+            return;
+        }
+        // SAFETY: the state is still installed, so the `BundleV2` owning it has
+        // not been torn down (see the type doc); the callback body's borrows of
+        // `*bv2` have ended by the time its locals drop.
+        unsafe { (*self.bv2).async_ast_alloc.exit() };
+    }
 }
 
 bun_core::declare_scope!(Bundle, visible);
@@ -268,6 +362,25 @@ impl<'a> BundleV2<'a> {
         }
     }
 
+    /// Dev server: the (uninstalled) state the bundle was set up under; see
+    /// [`AsyncAstAlloc`].
+    pub fn adopt_async_ast_state(&mut self, state: Box<AstAllocState>) {
+        debug_assert!(self.asynchronous);
+        debug_assert!(matches!(self.async_ast_alloc.0, AsyncAstState::Disabled));
+        self.async_ast_alloc.0 = AsyncAstState::Parked(state);
+    }
+
+    /// Must be the first local of every event loop callback that works on the
+    /// graph, so the state is still installed when a later guard or the body
+    /// completes the bundle.
+    pub(crate) fn enter_async_ast_scope(&mut self) -> AsyncAstScope {
+        let installed = self.async_ast_alloc.enter(self.graph.heap.heap_ptr());
+        AsyncAstScope {
+            bv2: std::ptr::from_mut::<BundleV2<'a>>(self).cast::<BundleV2<'static>>(),
+            installed,
+        }
+    }
+
     // draft `on_parse_task_complete` / `deinit_without_freeing_arena`
     // removed — canonical bodies live in the later impl blocks below.
 }
@@ -2853,6 +2966,7 @@ pub mod bv2_impl {
                 asynchronous: false,
                 has_any_top_level_await_modules: false,
                 requested_exports: Vec::new(),
+                async_ast_alloc: super::AsyncAstAlloc::default(),
             });
             if let Some(bo) = bake_options {
                 // SAFETY: `bo.client_transpiler` is the caller's live, write-capable
@@ -4408,6 +4522,7 @@ pub mod bv2_impl {
 
     impl<'a> BundleV2<'a> {
         pub(crate) fn on_load(load: &mut jsc_api::JSBundler::Load, this: &mut BundleV2) {
+            let _ast_alloc = this.enter_async_ast_scope();
             this.graph.outstanding_loads.unlink(load);
             load.deferred = false;
             // `Load` is arena-allocated (no Drop); free its owned heap fields on every exit path.
@@ -4607,6 +4722,8 @@ pub mod bv2_impl {
 
     impl<'a> BundleV2<'a> {
         pub(crate) fn on_resolve(resolve: &mut jsc_api::JSBundler::Resolve, this: &mut BundleV2) {
+            // Declared before `_dec_guard`, whose drop may complete the bundle.
+            let _ast_alloc = this.enter_async_ast_scope();
             this.graph.outstanding_resolves.unlink(resolve);
             // RAII guard captures `this`
             // as a raw pointer so it does not hold a unique borrow across the body.
@@ -5022,6 +5139,10 @@ pub mod bv2_impl {
             for free in self.free_list.drain(..) {
                 drop(free);
             }
+
+            // The dev server tears a bundle down from inside the callback that
+            // completed it, so the state is still installed here.
+            self.async_ast_alloc.exit();
         }
 
         pub fn run_from_js_in_new_thread(
@@ -6917,6 +7038,7 @@ pub mod bv2_impl {
 
     impl<'a> BundleV2<'a> {
         pub fn on_notify_defer(&mut self) {
+            let _ast_alloc = self.enter_async_ast_scope();
             self.thread_lock.assert_locked();
             self.graph.deferred_pending += 1;
             self.decrement_scan_counter();
@@ -6931,6 +7053,7 @@ pub mod bv2_impl {
             parse_result: &mut parse_task::Result,
             this: &mut BundleV2,
         ) {
+            let _ast_alloc = this.enter_async_ast_scope();
             let _trace = crate::perf::trace("Bundler.onParseTaskComplete");
             // Borrowck rejects holding a `&this.graph` alias
             // across the `this.*` method calls below (each takes

src/runtime/bake/DevServer.rs

@@ -238,9 +238,6 @@ pub struct CurrentBundle {
     /// Owns the arena that `bv2.graph.heap` borrows (`'static` self-ref via the
     /// boxed allocation's stable address; same erasure as `bv2` above).
     pub heap: Box<bun_alloc::MimallocArena>,
-    /// Backs the small `AstVec`s built during bundle setup
-    /// (`start_async_bundle`'s AST scope); dropped with the bundle.
-    pub ast_alloc_state: Option<Box<bun_alloc::ast_alloc::AstAllocState>>,
     /// Information BundleV2 needs to finalize the bundle
     pub(crate) start_data: bundler::bundle_v2::DevServerInput,
     /// Started when the bundle was queued
@@ -3169,8 +3166,8 @@ impl DevServer {
         let heap: Box<bun_alloc::MimallocArena> = Box::new(bun_alloc::MimallocArena::new());
         // Borrows `heap`, so AST nodes built during bundle setup
         // live exactly as long as the bundle. The arena-allocated allocator
-        // never runs `Drop`; the `AstAllocState` is taken into `CurrentBundle`
-        // on success and recycled by the guard below on error paths.
+        // never runs `Drop`; the `AstAllocState` is handed to `bv2` on success
+        // and recycled by the guard below on error paths.
         let ast_memory_store: *mut bun_ast::ASTMemoryAllocator =
             heap.alloc(bun_ast::ASTMemoryAllocator::borrowing(&heap));
         struct ReleaseAstState(*mut bun_ast::ASTMemoryAllocator);
@@ -3262,16 +3259,18 @@ impl DevServer {
             bt
         })?;
         drop(entry_points);
-        // End the AST scope and move its state into the bundle so the small
-        // `AstVec`s built during setup stay alive until the bundle completes.
+        // End the AST scope; the bundle reinstalls its state for the callbacks
+        // that finish the bundle (`BundleV2::enter_async_ast_scope`).
         drop(ast_scope);
         // SAFETY: `ast_memory_store` lives in `heap`; the scope above has
         // exited, so no `&mut` to the allocator is live.
         let ast_alloc_state = unsafe { (*ast_memory_store).take_ast_state() };
+        bv2.adopt_async_ast_state(
+            ast_alloc_state.unwrap_or_else(bun_alloc::ast_alloc::acquire_state),
+        );
         self.current_bundle = Some(CurrentBundle {
             bv2,
             heap,
-            ast_alloc_state,
             timer,
             start_data,
             had_reload_event,

test/bake/dev-server-memory.test.ts

@@ -0,0 +1,129 @@
+import { expect, test } from "bun:test";
+import { bunEnv, bunExe, tempDir } from "harness";
+import { join } from "node:path";
+
+// The graph work that finishes a dev server bundle (parse completions, linking)
+// runs on the JS thread. The `AstAlloc`-backed structures it builds there (among
+// them one entry per export of every re-bundled module) used to land on the
+// global mimalloc heap, where nothing ever frees them, so every rebuild leaked
+// about one block per export. The fixture reports the live block count of that
+// heap (`heapStats({ dump: true })`, heap seq 0), which, unlike RSS, is exact:
+// once the bundle's allocations die with the bundle it stays flat.
+const EXPORTS = 2000;
+const WARMUP_REBUILDS = 3;
+const MEASURED_REBUILDS = 15;
+
+function moduleSource(revision: number) {
+  let source = `export const revision = ${revision};\n`;
+  for (let i = 0; i < EXPORTS; i++) {
+    source += `export const e${i} = ${i};\n`;
+  }
+  return source;
+}
+
+test("rebuilding a module does not leak the bundle's allocations", async () => {
+  using dir = tempDir("dev-server-memory", {
+    "index.html": `<!doctype html><script type="module" src="./script.ts"></script>`,
+    "script.ts": `import * as mod from "./mod.ts";\nconsole.log(Object.keys(mod).length);\n`,
+    "mod.ts": moduleSource(0),
+    "server.ts": `
+      import { heapStats } from "bun:jsc";
+      import index from "./index.html";
+
+      const server = Bun.serve({
+        port: 0,
+        development: true,
+        routes: {
+          "/": index,
+          "/live-blocks": () => {
+            Bun.gc(true);
+            const mainHeap = heapStats({ dump: true }).mimallocDump.heaps.find(heap => heap.seq === 0);
+            return new Response(String(mainHeap.pages.reduce((blocks, page) => blocks + page.used, 0)));
+          },
+        },
+      });
+      console.log(server.port);
+    `,
+  });
+
+  await using proc = Bun.spawn({
+    cmd: [bunExe(), "server.ts"],
+    cwd: String(dir),
+    env: bunEnv,
+    stdout: "pipe",
+    stderr: "pipe",
+  });
+
+  // The dev server prints "Reloaded in <n>ms" once per completed rebuild.
+  let stderr = "";
+  let reloads = 0;
+  let exited: Error | undefined;
+  let waiter: { target: number; resolve(): void; reject(err: Error): void } | undefined;
+  const stderrClosed = (async () => {
+    const decoder = new TextDecoder();
+    for await (const chunk of proc.stderr) {
+      stderr += decoder.decode(chunk, { stream: true });
+      reloads = stderr.match(/Reloaded in /g)?.length ?? 0;
+      if (waiter && reloads >= waiter.target) {
+        waiter.resolve();
+        waiter = undefined;
+      }
+    }
+    exited = new Error(`dev server exited after ${reloads} reloads:\n${stderr}`);
+    waiter?.reject(exited);
+  })();
+
+  const stdout = proc.stdout.getReader();
+  let firstLine = "";
+  while (!firstLine.includes("\n")) {
+    const { value, done } = await stdout.read();
+    if (done) {
+      await stderrClosed;
+      throw exited;
+    }
+    firstLine += Buffer.from(value).toString();
+  }
+  const origin = `http://localhost:${Number.parseInt(firstLine, 10)}`;
+
+  async function get(path: string) {
+    const response = await fetch(origin + path);
+    const body = await response.text();
+    expect(response.status).toBe(200);
+    return body;
+  }
+
+  async function rebuild(revision: number) {
+    // An exit that happened while nothing was waiting is reported here; one
+    // that happens while waiting rejects the waiter.
+    if (exited) throw exited;
+    const { promise, resolve, reject } = Promise.withResolvers<void>();
+    waiter = { target: reloads + 1, resolve, reject };
+    await Bun.write(join(String(dir), "mod.ts"), moduleSource(revision));
+    await promise;
+  }
+
+  async function liveBlocks(): Promise<number> {
+    while (true) {
+      const reloadsBefore = reloads;
+      // Page requests are answered once any in-flight bundle has finished.
+      await get("/");
+      const blocks = Number(await get("/live-blocks"));
+      // A write can surface as more than one watcher event; if an extra
+      // rebuild landed while sampling, sample again.
+      if (reloads === reloadsBefore) return blocks;
+    }
+  }
+
+  let revision = 0;
+  await get("/");
+  for (let i = 0; i < WARMUP_REBUILDS; i++) await rebuild(++revision);
+  const before = await liveBlocks();
+  for (let i = 0; i < MEASURED_REBUILDS; i++) await rebuild(++revision);
+  const after = await liveBlocks();
+
+  // Unfixed: at least EXPORTS blocks per rebuild. Fixed: a few hundred at most.
+  expect(after - before).toBeLessThan((EXPORTS * MEASURED_REBUILDS) / 4);
+
+  proc.kill();
+  await Promise.all([proc.exited, stderrClosed]);
+}, 60_000);