bundler: free the ServerComponentParseTask after it generates its file

src/bundler/ServerComponentParseTask.rs

@@ -8,10 +8,10 @@ use std::fmt::Write as _;
 use bun_alloc::{AllocError as OOM, Arena}; // bumpalo::Bump re-export
 use bun_collections::VecExt;
 
-use bun_ast::{Loc, Log, Source};
+use bun_ast::{Index, Loc, Log, Source};
+use bun_paths::fs::Path as FsPath;
 use bun_threading::thread_pool::Task as ThreadPoolTask;
 
-use bun_ast::ast_result::NamedExports;
 use bun_ast::{B, Binding, E, G, S, Stmt, symbol};
 use bun_ast::{ExprNodeList, LocRef, StmtOrExpr, UseDirective};
 use bun_ast::{ImportKind, ImportRecordFlags};
@@ -24,19 +24,18 @@ use crate::cache::ExternalFreeFunction;
 use crate::options::{Loader, Target};
 use crate::parse_task::{self, ResultValue, Success, WatcherData, on_complete};
 
+/// One box per generated file; [`task_callback_wrap`] takes it back and frees it.
 pub(crate) struct ServerComponentParseTask {
     pub task: ThreadPoolTask,
     pub data: Data,
     // BACKREF (LIFETIMES.tsv) — written through in `on_complete`.
     // `ParentRef` (write-provenance via `NonNull::from(&mut self)` at construction)
     // so deref sites are safe; `None` only for the FRU `Default` placeholder.
     pub ctx: Option<bun_ptr::ParentRef<BundleV2<'static>, bun_ptr::Mut>>,
+    /// The generated file's own source record; moved into the result.
     pub source: Source,
 }
 
-// `ServerComponentParseTask` is bump-arena-allocated; boxing the large arm
-// would leak. The size diff is acceptable.
-#[allow(clippy::large_enum_variant)]
 pub enum Data {
     /// Generate server-side code for a "use client" module. Given the
     /// client ast, a "reference proxy" is created with identical exports.
@@ -46,18 +45,19 @@ pub enum Data {
 }
 
 pub struct ReferenceProxy {
-    pub(crate) other_source: Source,
-    pub(crate) named_exports: NamedExports,
+    /// The "use client" module the proxy stands in for.
+    pub(crate) client_path: FsPath<'static>,
+    pub(crate) client_source_index: Index,
+    /// In export order; bundle-arena copies (`BundleV2::copy_export_names_for_reference_proxy`).
+    pub(crate) export_names: &'static [&'static [u8]],
 }
 
 pub struct ClientEntryWrapper {
-    // Owned copy.
-    pub(crate) path: Box<[u8]>,
+    /// Bundle-arena or `'static` bytes; stored as-is in the `ImportRecord`.
+    pub(crate) path: &'static [u8],
 }
 
-/// Raw thread-pool callback. Recovers `&mut ServerComponentParseTask` from the
-/// intrusive `task` field and dispatches the parse, then posts the result back
-/// to the owning event loop.
+/// Thread-pool callback: takes the task box back, generates the file, frees it, posts the result.
 // CONCURRENCY: thread-pool callback — runs on worker threads, one task per
 // `ServerComponentParseTask` (heap-allocated, scheduled exactly once). Writes:
 // own fields + `Log` (local) + result is posted via
@@ -66,10 +66,17 @@ pub struct ClientEntryWrapper {
 // backref to a `Send` type and `Source`/`Data` payloads are bundle-arena
 // slices.
 fn task_callback_wrap(thread_pool_task: *mut ThreadPoolTask) {
-    // SAFETY: `thread_pool_task` points to the `task` field of a heap-allocated
-    // `ServerComponentParseTask` enqueued by BundleV2; offset_of recovers the parent.
-    let task: &mut ServerComponentParseTask = unsafe {
-        &mut *(bun_core::from_field_ptr!(ServerComponentParseTask, task, thread_pool_task))
+    // SAFETY: `thread_pool_task` is the `task` field of the
+    // `ServerComponentParseTask` that `enqueue_server_component_generated_file`
+    // leaked with `heap::into_raw` and scheduled; the pool runs a task exactly
+    // once and never touches it after this callback, so ownership of the box
+    // is ours to reclaim.
+    let mut task: Box<ServerComponentParseTask> = unsafe {
+        bun_core::heap::take(bun_core::from_field_ptr!(
+            ServerComponentParseTask,
+            task,
+            thread_pool_task
+        ))
     };
 
     // `ctx` is a `ParentRef` BACKREF to the owning BundleV2 (set at enqueue).
@@ -84,11 +91,13 @@ fn task_callback_wrap(thread_pool_task: *mut ThreadPoolTask) {
     // worker-owned bump arena; lives for the worker's lifetime.
     let arena: &Arena = worker.arena();
 
-    let value = match task_callback(task, &mut log, arena) {
+    let value = match task_callback(&mut task, &mut log, arena) {
         Ok(success) => ResultValue::Success(success),
         // Only possible error is OOM; abort like `bun.outOfMemory()`.
         Err(_oom) => bun_core::out_of_memory(),
     };
+    // Nothing in `value` refers to the task, so it is freed before the result is posted.
+    drop(task);
 
     let result = Box::new(parse_task::Result {
         // `ctx` already a `ParentRef<BundleV2>` with write provenance
@@ -215,21 +224,15 @@ impl Default for ServerComponentParseTask {
                 node: Default::default(),
                 callback: task_callback_wrap,
             },
-            data: Data::ClientEntryWrapper(ClientEntryWrapper {
-                path: Box::default(),
-            }),
+            data: Data::ClientEntryWrapper(ClientEntryWrapper { path: b"" }),
             ctx: None,
             source: Source::default(),
         }
     }
 }
 
 fn generate_client_entry_wrapper(data: &ClientEntryWrapper, b: &mut AstBuilder) -> Result<(), OOM> {
-    // `add_import_record` stores the slice raw in the `ImportRecord`; `data.path`
-    // outlives the bundle pass (owned by the heap-allocated task). Route through
-    // `StoreStr` so the lifetime erasure goes through one audited unsafe.
-    let path = bun_ast::StoreStr::new(&data.path[..]);
-    let record = b.add_import_record(path.slice(), ImportKind::Stmt)?;
+    let record = b.add_import_record(data.path, ImportKind::Stmt)?;
     let namespace_ref = b.new_symbol(symbol::Kind::Other, b"main")?;
     b.append_stmt(S::Import {
         namespace_ref,
@@ -257,8 +260,6 @@ fn generate_client_reference_proxy(
         // config must be non-null to enter this function
         .unwrap_or_else(|| unreachable!());
 
-    let client_named_exports = &data.named_exports;
-
     // `add_import_stmt` stores the slices raw in `ImportRecord`/`ClauseItem`s;
     // the framework config outlives the bundle pass. Route through `StoreStr`
     // so the lifetime erasure goes through one audited unsafe.
@@ -275,7 +276,7 @@ fn generate_client_reference_proxy(
         // that information is not yet available since chunks are not
         // computed. The unique_key replacement system is used here.
         if ctx.transpiler().options.has_dev_server() {
-            b.bump.alloc_slice_copy(data.other_source.path.pretty)
+            b.bump.alloc_slice_copy(data.client_path.pretty)
         } else {
             let mut buf = bun_alloc::ArenaString::new_in(b.bump);
             write!(
@@ -284,16 +285,15 @@ fn generate_client_reference_proxy(
                 crate::chunk::UniqueKey {
                     prefix: ctx.unique_key,
                     kind: crate::chunk::QueryKind::Scb,
-                    index: data.other_source.index.0,
+                    index: data.client_source_index.0,
                 },
             )
             .map_err(|_| OOM)?;
             buf.into_bump_str().as_bytes()
         },
     ));
 
-    for key in client_named_exports.keys() {
-        let key: &[u8] = key.as_ref();
+    for &key in data.export_names {
         let is_default = key == b"default";
 
         // This error message is taken from
@@ -309,7 +309,7 @@ fn generate_client_reference_proxy(
                         "client function from the server, it can only be rendered as a ",
                         "Component or passed to props of a Client Component.",
                     ),
-                    module_path = bstr::BStr::new(data.other_source.path.pretty),
+                    module_path = bstr::BStr::new(data.client_path.pretty),
                 )
             } else {
                 write!(
@@ -359,7 +359,7 @@ fn generate_client_reference_proxy(
                     ..Default::default()
                 }),
                 module_path,
-                b.new_expr(E::String::init(b.bump.alloc_slice_copy(key))),
+                b.new_expr(E::String::init(key)),
             ]),
             ..Default::default()
         });

src/bundler/bundle_v2.rs

@@ -3728,6 +3728,24 @@ pub mod bv2_impl {
             Ok(source_index.get())
         }
 
+        /// Workers may not read `graph.ast` while this thread appends to it, hence the copies.
+        fn copy_export_names_for_reference_proxy(
+            &self,
+            named_exports: &crate::bundled_ast::NamedExports,
+        ) -> &'static [&'static [u8]] {
+            // SAFETY: same contract as `interned_slice` — the arena outlives the
+            // bundle pass, and the proxy task (and the AST it produces, which
+            // stores these slices) is consumed within it.
+            let arena: &'static bun_alloc::Arena =
+                unsafe { bun_ptr::detach_lifetime_ref::<bun_alloc::Arena>(self.arena()) };
+            arena.alloc_slice_fill_iter(
+                named_exports
+                    .keys()
+                    .iter()
+                    .map(|name| -> &'static [u8] { arena.alloc_slice_copy(&name[..]) }),
+            )
+        }
+
         /// Enqueue a ServerComponentParseTask.
         /// `source_without_index` is copied and assigned a new source index. That index is returned.
         pub(crate) fn enqueue_server_component_generated_file(
@@ -3754,9 +3772,7 @@ pub mod bv2_impl {
             })?;
             let _ = self.graph.ast.append(JSAst::empty_in(self.graph.heap)); // OOM/capacity: fire-and-forget
 
-            // `bun.new(ServerComponentParseTask, …)` — heap-owned by the
-            // worker pool; freed via `bun.destroy` in `on_complete` after the
-            // result posts back to the bundle thread.
+            // Freed by the pool callback (`task_callback_wrap`).
             let task = bun_core::heap::into_raw(Box::new(ServerComponentParseTask {
                 data,
                 // SAFETY: `from_mut(self)` is the live bundle (write provenance for
@@ -3774,7 +3790,9 @@ pub mod bv2_impl {
 
             self.increment_scan_counter();
 
-            // SAFETY: `task` is the just-allocated arena box; sole reference here.
+            // SAFETY: `task` is the just-allocated box and this is its only
+            // pointer; projecting through it (not through a `&mut` to the field)
+            // keeps whole-allocation provenance for the callback's `heap::take`.
             self.graph
                 .pool()
                 .worker_pool()
@@ -7176,11 +7194,8 @@ pub mod bv2_impl {
                     }
                     result.ast.import_records = import_records;
 
-                    // `result.ast` is moved into `graph.ast` and `result.source` was
-                    // swapped earlier, so snapshot the data the use-directive block
-                    // needs *before* the move. Only paid for files that hit the SCB gate.
-                    let named_exports_for_scb = if result.use_directive != crate::UseDirective::None
-                        && {
+                    let is_server_component_boundary =
+                        result.use_directive != crate::UseDirective::None && {
                             let separate = this
                                 .framework
                                 .as_ref()
@@ -7196,11 +7211,7 @@ pub mod bv2_impl {
                             } else {
                                 is_client != is_browser
                             }
-                        } {
-                        Some(result.ast.named_exports.clone().expect("oom"))
-                    } else {
-                        None
-                    };
+                        };
 
                     let result_heap = *result.ast.parts.allocator();
                     this.graph.ast.set(
@@ -7219,7 +7230,8 @@ pub mod bv2_impl {
                         .expect("oom");
                     }
 
-                    if let Some(named_exports) = named_exports_for_scb {
+                    // Index the boundary and enqueue the files for its other side.
+                    if is_server_component_boundary {
                         if result.use_directive == crate::UseDirective::Server {
                             bun_core::todo_panic!("\"use server\"");
                         }
@@ -7241,17 +7253,21 @@ pub mod bv2_impl {
 
                         let (reference_source_index, ssr_index) = if separate_ssr_graph {
                             // Enqueue two files, one in server graph, one in ssr graph.
-                            let other_source =
-                                this.graph.input_files.items_source()[result_source_index].clone();
-                            let scb_source =
-                                this.graph.input_files.items_source()[result_source_index].clone();
+                            let export_names = this.copy_export_names_for_reference_proxy(
+                                &this.graph.ast.items_named_exports()[result_source_index],
+                            );
+                            let client_source =
+                                &this.graph.input_files.items_source()[result_source_index];
+                            let proxy = crate::ServerComponentParseTask::ReferenceProxy {
+                                client_path: client_source.path,
+                                client_source_index: client_source.index,
+                                export_names,
+                            };
+                            let scb_source = client_source.clone();
                             let reference_source_index = this
                                 .enqueue_server_component_generated_file(
                                     crate::ServerComponentParseTask::Data::ClientReferenceProxy(
-                                        crate::ServerComponentParseTask::ReferenceProxy {
-                                            other_source,
-                                            named_exports,
-                                        },
+                                        proxy,
                                     ),
                                     scb_source,
                                 )

test/bake/deinitialization.test.ts

@@ -1,4 +1,5 @@
-import { bunEnv, bunExe } from "harness";
+import { expect, test } from "bun:test";
+import { bunEnv, bunExe, isASAN, tempDir } from "harness";
 import path from "node:path";
 
 test("dev server deinitializes itself", () => {
@@ -13,3 +14,96 @@ test("dev server deinitializes itself", () => {
   // The child runs a whole `bun test` suite (nine GC-heavy cases plus leak
   // reporting at exit), which takes longer than the 5s default under ASAN.
 }, 60_000);
+
+// With separateSSRGraph, every "use client" module the server graph imports
+// makes the bundler generate a "reference proxy" module in its place. The
+// proxy is built on the thread pool from a heap-allocated
+// ServerComponentParseTask, which used to be scheduled and never freed (one
+// per client module per bundle). LeakSanitizer reports the task at exit, so
+// this needs the ASAN build. The rendered output checks what the task hands
+// to the generator: the client module's path and its export names, which the
+// proxy passes to registerClientReference.
+test.skipIf(!isASAN)(
+  'bundling a "use client" reference proxy does not leak its ServerComponentParseTask',
+  async () => {
+    using dir = tempDir("bake-reference-proxy-leak", {
+      "main.ts": `
+        using server = Bun.serve({
+          port: 0,
+          development: true,
+          app: {
+            framework: {
+              serverComponents: {
+                separateSSRGraph: true,
+                serverRuntimeImportSource: "./framework/server.ts",
+                serverRegisterClientReferenceExport: "registerClientReference",
+              },
+              fileSystemRouterTypes: [
+                {
+                  root: "routes",
+                  serverEntryPoint: "./framework/server.ts",
+                  style: "nextjs-pages",
+                },
+              ],
+            },
+          },
+        });
+        const response = await fetch(server.url);
+        console.log(response.status, await response.text());
+      `,
+      "framework/server.ts": `
+        export function render(request, meta) {
+          return new Response(meta.pageModule.default());
+        }
+        // Called once per export of the client module by the generated proxy.
+        export function registerClientReference(value, file, exportName) {
+          return () => file + "#" + exportName;
+        }
+      `,
+      // The server graph imports both client modules, so it gets a proxy for
+      // each: the proxies, not the modules, run on the server, which is why
+      // Empty's side effect must not be visible to the route.
+      "routes/index.ts": `
+        import Widget, { Alpha, Beta } from "../components/Widget";
+        import "../components/Empty";
+        export default () => [Widget(), Alpha(), Beta(), String(globalThis.emptyRanOnServer)].join(" ");
+      `,
+      "components/Widget.ts": `
+        "use client";
+        export function Alpha() {}
+        export function Beta() {}
+        export default function Widget() {}
+      `,
+      // A client module with no exports gets a proxy with no exports.
+      "components/Empty.ts": `
+        "use client";
+        globalThis.emptyRanOnServer = true;
+      `,
+    });
+
+    await using proc = Bun.spawn({
+      cmd: [bunExe(), "main.ts"],
+      cwd: String(dir),
+      env: {
+        ...bunEnv,
+        BUN_DESTRUCT_VM_ON_EXIT: "1",
+        ASAN_OPTIONS: [bunEnv.ASAN_OPTIONS, "detect_leaks=1"].filter(Boolean).join(":"),
+        LSAN_OPTIONS: `print_suppressions=0:suppressions=${path.join(import.meta.dir, "../leaksan.supp")}`,
+      },
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+
+    expect(stdout).toBe(
+      "200 components/Widget.ts#default components/Widget.ts#Alpha components/Widget.ts#Beta undefined\n",
+    );
+    // The dev server logs its bundle timing here; a leak report would follow it.
+    expect(stderr).not.toContain("LeakSanitizer");
+    expect(exitCode).toBe(0);
+  },
+  // Starting a dev server and bundling the route takes a few seconds under
+  // ASAN, and when LSan does find something, symbolizing the report against
+  // the debug binary takes longer still.
+  90_000,
+);