diff --git a/src/bundler/bundle_v2.rs b/src/bundler/bundle_v2.rs index 24999fd82b68..b647af31b388 100644 --- a/src/bundler/bundle_v2.rs +++ b/src/bundler/bundle_v2.rs @@ -4,8 +4,11 @@ // bodies live in the `bv2_impl` module below. // ══════════════════════════════════════════════════════════════════════════ +use core::cell::Cell; use core::ptr::NonNull; +use std::rc::Rc; +use bun_alloc::ast_alloc::{self, AstAllocState}; use bun_collections::{ArrayHashMap, StringHashMap}; use bun_core::ThreadLock; @@ -140,6 +143,104 @@ 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>, + + /// Declared last: `graph` / `linker` hold `AstVec`s backed by the parked + /// state's inline chunk, so it must drop after them. + pub(crate) async_ast_alloc: Rc, +} + +/// 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`]). +/// +/// Shared (`Rc`) with the guard of the callback that completes the bundle, +/// since that frees the `BundleV2` before the guard drops. +#[derive(Default)] +pub(crate) struct AsyncAstAlloc(Cell); + +#[derive(Default)] +enum AsyncAstState { + /// Synchronous bundle: the owning thread keeps its own state installed. + #[default] + Disabled, + Parked(Box), + Installed { + id: *const AstAllocState, + /// Thread-local occupant displaced by `enter`, reinstated by `exit`. + displaced: Option>, + }, +} + +impl AsyncAstAlloc { + fn adopt(&self, state: Box) { + let previous = self.0.replace(AsyncAstState::Parked(state)); + debug_assert!(matches!(previous, AsyncAstState::Disabled)); + } + + /// Installs the parked state; `false` (and a no-op) unless one is parked. + fn enter(&self, spill: *mut bun_alloc::mimalloc::Heap) -> bool { + let mut state = match self.0.take() { + AsyncAstState::Parked(state) => state, + other => { + self.0.set(other); + return false; + } + }; + state.set_spill_heap(spill); + let id: *const AstAllocState = &raw const *state; + self.0.set(AsyncAstState::Installed { + id, + displaced: ast_alloc::swap_state(Some(state)), + }); + true + } + + /// Uninstalls and parks the state again; no-op unless it is installed and + /// still the thread's active state. + fn exit(&self) { + let (id, displaced) = match self.0.take() { + AsyncAstState::Installed { id, displaced } => (id, displaced), + other => { + self.0.set(other); + return; + } + }; + if !core::ptr::eq(ast_alloc::active_state_id(), id) { + debug_assert!( + false, + "AsyncAstAlloc::exit: another AstAllocState is installed on top of the bundle's" + ); + self.0.set(AsyncAstState::Installed { id, displaced }); + return; + } + let state = ast_alloc::swap_state(displaced) + .expect("AsyncAstAlloc::exit: the active state's identity matched"); + self.0.set(AsyncAstState::Parked(state)); + } +} + +impl Drop for AsyncAstAlloc { + fn drop(&mut self) { + self.exit(); + if let AsyncAstState::Parked(state) = self.0.take() { + ast_alloc::release_state(state); + } + } +} + +/// Returned by [`BundleV2::enter_async_ast_scope`]; `Some` only for the guard +/// that installed the state, which uninstalls it on drop. +#[must_use = "the state is uninstalled as soon as the guard drops"] +pub(crate) struct AsyncAstScope(Option>); + +impl Drop for AsyncAstScope { + fn drop(&mut self) { + if let Some(alloc) = &self.0 { + alloc.exit(); + } + } } bun_core::declare_scope!(Bundle, visible); @@ -268,6 +369,24 @@ 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) { + debug_assert!(self.asynchronous); + self.async_ast_alloc.adopt(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(&self) -> AsyncAstScope { + AsyncAstScope( + self.async_ast_alloc + .enter(self.graph.heap.heap_ptr()) + .then(|| Rc::clone(&self.async_ast_alloc)), + ) + } + // draft `on_parse_task_complete` / `deinit_without_freeing_arena` // removed — canonical bodies live in the later impl blocks below. } @@ -2853,6 +2972,7 @@ pub mod bv2_impl { asynchronous: false, has_any_top_level_await_modules: false, requested_exports: Vec::new(), + async_ast_alloc: Default::default(), }); if let Some(bo) = bake_options { // SAFETY: `bo.client_transpiler` is the caller's live, write-capable @@ -4408,6 +4528,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 +4728,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 +5145,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 +7044,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 +7059,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 diff --git a/src/runtime/bake/DevServer.rs b/src/runtime/bake/DevServer.rs index e8f6e34786c5..327991fea490 100644 --- a/src/runtime/bake/DevServer.rs +++ b/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, - /// Backs the small `AstVec`s built during bundle setup - /// (`start_async_bundle`'s AST scope); dropped with the bundle. - pub ast_alloc_state: Option>, /// 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 = 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, diff --git a/test/bake/dev-server-memory.test.ts b/test/bake/dev-server-memory.test.ts new file mode 100644 index 000000000000..c5b0c0e0e76b --- /dev/null +++ b/test/bake/dev-server-memory.test.ts @@ -0,0 +1,131 @@ +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": ``, + "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); + if (!mainHeap) throw new Error("heapStats dump has no heap with 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 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) { + if (exited) throw exited; + const response = await fetch(origin + path); + const body = await response.text(); + if (response.status !== 200) throw new Error(`GET ${path} responded with ${response.status}:\n${body}`); + return body; + } + + async function rebuild(revision: number) { + // An exit 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(); + waiter = { target: reloads + 1, resolve, reject }; + await Bun.write(join(String(dir), "mod.ts"), moduleSource(revision)); + await promise; + } + + async function liveBlocks(): Promise { + 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);