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,125 @@ 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>>,
+
+    /// See [`AsyncAstAlloc`]. Declared last: the state's inline chunk backs
+    /// small `AstVec`s stored in `graph` / `linker`, so it must outlive their
+    /// drop glue (struct fields drop in declaration order).
+    pub(crate) async_ast_alloc: AsyncAstAlloc,
+}
+
+/// The `AstAlloc` state of an `asynchronous` (dev server) bundle.
+///
+/// `AstAlloc` allocates through the `AstAllocState` installed on the calling
+/// thread and leaks on the global heap when there is none (`deallocate` is a
+/// no-op). A synchronous bundle (`Bun.build`, the CLI) keeps one installed on
+/// its own thread for the whole pass. A dev server bundle only has one while
+/// `DevServer::start_async_bundle` sets it up; the rest of the graph work
+/// arrives later as JS event loop callbacks (`on_parse_task_complete`,
+/// `on_load`, `on_resolve`, `on_notify_defer`, and `finish_from_bake_dev_server`
+/// reached from them) with nothing installed, so everything those built through
+/// `AstAlloc` (`LinkerGraph::load`'s per-file `resolved_exports`, `clone_ast`,
+/// `InputFile::additional_files`, ...) leaked once per rebuild.
+///
+/// The setup state, which spills into `graph.heap`, is parked here between
+/// callbacks and installed for the duration of each one, so those allocations
+/// die with the bundle heap like everything else the bundle owns.
+#[derive(Default)]
+pub(crate) struct AsyncAstAlloc(AsyncAstState);
+
+#[derive(Default)]
+enum AsyncAstState {
+    /// Synchronous bundle: [`BundleV2::enter_async_ast_scope`] is a no-op.
+    #[default]
+    Disabled,
+    /// Between callbacks.
+    Parked(Box<AstAllocState>),
+    /// During a callback. A nested `enter` is a no-op.
+    Installed {
+        /// Identity of the installed box.
+        id: *const AstAllocState,
+        /// The thread-local occupant `enter` displaced; reinstated by `exit`.
+        displaced: Option<Box<AstAllocState>>,
+    },
+}
+
+impl AsyncAstAlloc {
+    /// Install the parked state, spilling into `spill`. Returns the installed
+    /// box's identity, or null when nothing was installed.
+    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
+    }
+
+    /// Uninstall the state, reinstating whatever `enter` displaced, and park it
+    /// 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) {
+        // Normally already parked by `deinit_without_freeing_arena`; this keeps
+        // the thread-local from pointing at the freed box otherwise.
+        self.exit();
+    }
+}
+
+/// Returned by [`BundleV2::enter_async_ast_scope`]; uninstalls the bundle's
+/// state when dropped.
+///
+/// The callback holding this guard may be the one that completes the bundle:
+/// `finish_from_bake_dev_server` → `DevServer::finalize_bundle` frees the
+/// `BundleV2` before returning into the callback. That teardown uninstalls the
+/// state itself (`deinit_without_freeing_arena`, then `AsyncAstAlloc::drop`),
+/// so this guard dereferences `bv2` only while the state it installed is still
+/// the active one, which proves the `BundleV2` is still alive.
+#[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 this guard installed is still active, so the
+        // `BundleV2` owning it has not been torn down (see the type doc). The
+        // guard is a local of a callback running on the bundle's own thread,
+        // and every borrow of `*bv2` taken by the callback body has ended by
+        // the time its locals drop.
+        unsafe { (*self.bv2).async_ast_alloc.exit() };
+    }
 }
 
 bun_core::declare_scope!(Bundle, visible);
@@ -268,6 +388,28 @@ impl<'a> BundleV2<'a> {
         }
     }
 
+    /// Dev server: hand over the `AstAllocState` the bundle was set up under,
+    /// once setup's own scope has exited, so the event loop callbacks that
+    /// finish the bundle allocate through it as well (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);
+    }
+
+    /// Install the bundle's parked `AstAllocState`, if it has one, until the
+    /// returned guard drops. Every JS event loop callback that works on the
+    /// graph declares this as its first local, so the state is still installed
+    /// while anything declared later (or called from 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 +2995,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 +4551,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 +4751,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 +5168,11 @@ pub mod bv2_impl {
             for free in self.free_list.drain(..) {
                 drop(free);
             }
+
+            // A dev server bundle is torn down from inside the callback that
+            // completed it, i.e. while its state is installed; the heap that
+            // state spills into is destroyed right after this returns.
+            self.async_ast_alloc.exit();
         }
 
         pub fn run_from_js_in_new_thread(
@@ -6917,6 +7068,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 +7083,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,21 @@ 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 and hand its state to the bundle: the small
+        // `AstVec`s built during setup live in it, and the event loop callbacks
+        // that finish the bundle (`BundleV2::enter_async_ast_scope`) reinstall
+        // it so their `AstAlloc` allocations land in `heap` too instead of
+        // leaking on the global heap once per rebuild.
         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,121 @@
+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 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;
+      }
+    }
+    waiter?.reject(new Error(`dev server exited after ${reloads} reloads:\n${stderr}`));
+  })();
+
+  const stdout = proc.stdout.getReader();
+  let firstLine = "";
+  while (!firstLine.includes("\n")) {
+    const { value, done } = await stdout.read();
+    if (done) throw new Error(`dev server exited before printing its port:\n${stderr}`);
+    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) {
+    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);