bundler: join the work still on the pools before the bundle is torn down

src/ast/ast_memory_allocator.rs

@@ -267,7 +267,7 @@ impl ASTMemoryAllocator {
     /// Per-iteration reset for hot reuse paths (`initialize_mini_store`'s
     /// per-workspace-child re-entry). Thin delegate to
     /// [`bun_alloc::Arena::reset_retain_with_limit`]; the cold init paths
-    /// (`bundler::ThreadPool::Worker::init`, `BundleThread::generate_in_new_
+    /// (`bundler::ThreadPool::Worker::create`, `BundleThread::generate_in_new_
     /// thread`) keep calling [`Self::reset`].
     pub fn reset_retain_with_limit(&mut self, limit: usize) {
         if self.arena_dirty {

src/bundler/BundleThread.rs

@@ -296,9 +296,9 @@ impl<C: CompletionStruct> BundleThread<C> {
 
         // Straight-line teardown: log copy
         // runs on both paths; `completeOnBundleThread` only on success (the error
-        // path's `set_result(Err)` + complete happens in `thread_main`). The
-        // `deinitWithoutFreeingArena` + wait-group drain live inside `init_and_run`
-        // (it owns `this`).
+        // path's `set_result(Err)` + complete happens in `thread_main`).
+        // `deinit_without_freeing_arena` lives inside `init_and_run` (it owns
+        // `this`).
         let mut out_log = bun_ast::Log::init();
         // SAFETY: `transpiler.log` is the arena-allocated `*mut Log` set up by
         // `configure_bundler`; valid for the lifetime of `heap`. Raw deref so the

src/bundler/ParseTask.rs

@@ -362,7 +362,13 @@ unsafe fn io_task_callback(task: *mut ThreadPoolLib::Task) {
     // provenance covers the full `ParseTask` and the `&mut` is unique per the
     // CONCURRENCY note above.
     let parse_task = unsafe { &mut *bun_core::from_field_ptr!(ParseTask, io_task, task) };
+    // Captured before the run: after it, a worker thread may own `*parse_task`.
+    // SAFETY: teardown waits for the count this callback holds
+    // (`ThreadPool::wait_for_io_tasks`), so the bundle is live until `finish_io_task`.
+    let pool: *const crate::ThreadPool = unsafe { parse_task.ctx() }.graph.pool();
     parse_worker::run_from_thread_pool(parse_task);
+    // SAFETY: `pool` scheduled this task; this is the last access to the bundle.
+    unsafe { crate::ThreadPool::finish_io_task(pool) };
 }
 
 // CONCURRENCY: see `io_task_callback` — same task, different intrusive field.

src/bundler/ThreadPool.rs

@@ -15,7 +15,7 @@ use bun_alloc::Arena as ThreadLocalArena;
 use bun_collections::{ArrayHashMap, MapEntry};
 use bun_core::{self, env_var, output as Output};
 use bun_sys::Fd;
-use bun_threading::{Mutex, thread_pool as ThreadPoolLib};
+use bun_threading::{Mutex, WaitGroup, thread_pool as ThreadPoolLib};
 
 use crate::cache::{Contents, Entry as CacheEntry};
 use crate::linker_context_mod::StmtList;
@@ -49,6 +49,11 @@ pub struct ThreadPool {
     // `wake_for_idle_events`) take `&self` — so the safe `Deref` projection is
     // sufficient and the per-read `unsafe { p.as_ref() }` disappears.
     pub(crate) io_pool: Option<bun_ptr::ParentRef<ThreadPoolLib::ThreadPool>>,
+    /// One count per task on `io_pool`, finished when its callback returns.
+    /// That callback ends by scheduling the task onto `worker_pool`, and the
+    /// parse can complete (`graph.pending_items` reaching zero) before it has
+    /// returned from that, so teardown joins this before it frees `worker_pool`.
+    io_tasks_in_flight: WaitGroup,
     // Conditionally owned via `worker_pool_is_owned`; kept raw so callers
     // (bundle_v2.rs) can dereference for `wake_for_idle_events()` without a
     // borrow on `ThreadPool`.
@@ -208,6 +213,7 @@ impl ThreadPool {
             } else {
                 None
             },
+            io_tasks_in_flight: WaitGroup::init(),
             // BACKREF: lifetime erased behind the raw pointer.
             v2: std::ptr::from_ref(v2).cast::<BundleV2<'static>>(),
             worker_pool_is_owned: false,
@@ -216,28 +222,50 @@ impl ThreadPool {
         }
     }
 
+    /// Blocks until every IO callback has returned. `&self` on purpose: IO
+    /// threads still read the pool until this returns.
+    pub(crate) fn wait_for_io_tasks(&self) {
+        self.io_tasks_in_flight.wait();
+    }
+
+    /// The last step of `io_task_callback`.
+    ///
+    /// # Safety
+    /// `this` scheduled the task. It may be freed once [`Self::wait_for_io_tasks`]
+    /// returns, which this call permits, so nothing may touch it afterwards.
+    pub(crate) unsafe fn finish_io_task(this: *const Self) {
+        // SAFETY: caller contract; `finish_raw` is the matching last access.
+        unsafe { WaitGroup::finish_raw(&raw const (*this).io_tasks_in_flight) };
+    }
+
     /// Explicit teardown (no Drop on
     /// `ThreadPool` because `Graph.pool` is `NonNull<ThreadPool>` and the arena
-    /// owns the storage).
+    /// owns the storage). Only valid once no task of this pool is running:
+    /// the caller drains the parse tasks and calls [`Self::wait_for_io_tasks`].
     pub(crate) fn deinit(&mut self) {
         if self.worker_pool_is_owned {
             // SAFETY: worker_pool was heap-allocated in `init()` when owned.
             unsafe { drop(bun_core::heap::take(self.worker_pool)) };
-            self.worker_pool = ptr::null_mut();
         }
+        // Also for a shared pool, so that `worker_pool()` catches a late task on every path.
+        self.worker_pool = ptr::null_mut();
         if Self::uses_io_pool() {
             io_thread_pool::release();
         }
     }
 
-    /// Safe accessor for the underlying `bun_threading::ThreadPool`. The
-    /// pointer is set in `init`/`init_with_pool` and never null while `self`
-    /// is observable; encapsulating the deref keeps callers out of `unsafe`.
+    /// Safe accessor for the underlying `bun_threading::ThreadPool`. Set in
+    /// `init`/`init_with_pool`, nulled by `deinit`. `self` stays readable in the
+    /// arena after `deinit`, so a task that runs too late fails here by name.
     #[inline]
     pub(crate) fn worker_pool(&self) -> &ThreadPoolLib::ThreadPool {
-        debug_assert!(!self.worker_pool.is_null());
-        // SAFETY: `worker_pool` is initialized before any caller can observe
-        // `self` and lives until `deinit`; all driver methods take `&self`.
+        assert!(
+            !self.worker_pool.is_null(),
+            "bundler worker pool used after ThreadPool::deinit"
+        );
+        // SAFETY: non-null means `deinit` has not run. Until then an owned pool is
+        // not freed and a shared one is the process-lifetime `WorkPool`; all
+        // driver methods take `&self`.
         unsafe { &*self.worker_pool }
     }
 
@@ -343,6 +371,8 @@ impl ThreadPool {
                     ParseTaskStage::NeedsSourceCode => {
                         // io_pool is Some when uses_io_pool().
                         let io = self.io_pool_ref().unwrap();
+                        // Finished by `finish_io_task` at the end of `io_task_callback`.
+                        self.io_tasks_in_flight.add_one();
                         schedule_fn(
                             io,
                             ThreadPoolLib::Batch::from(&raw mut (*parse_task).io_task),
@@ -401,40 +431,17 @@ impl ThreadPool {
 
     #[cold]
     fn get_worker_slow(&self, id: ThreadId) -> &'static mut Worker {
-        let worker: *mut Worker;
-        {
-            let mut map = self.workers_assignments.lock();
-            match map.entry(id) {
-                MapEntry::Occupied(o) => {
-                    let w = *o.into_mut();
-                    drop(map);
-                    TLS_WORKER.set((self.generation, w));
-                    // SAFETY: map only stores live heap-allocated Workers (inserted below).
-                    return unsafe { &mut *w };
-                }
-                MapEntry::Vacant(v) => {
-                    // Allocate raw uninitialized storage; fully written via
-                    // `worker.write(...)` below before any read. Keep it as
-                    // `*mut MaybeUninit<Worker>` → `.cast()` instead of
-                    // `assume_init()` so we never materialize a `Box<Worker>`
-                    // whose payload violates `Worker`'s validity invariants
-                    // (niche-optimized `Option<_>` discriminants, the non-null
-                    // fn-pointer in `deinit_task.callback`, `bool` fields).
-                    worker = bun_core::heap::into_raw(Box::<Worker>::new_uninit()).cast::<Worker>();
-                    v.insert(worker);
-                }
-            }
-        }
-
-        // SAFETY: `worker` is freshly heap-allocated and exclusive on this
-        // thread until published via the map (already inserted above, but no
-        // other thread looks it up under a different `id`).
-        unsafe {
-            worker.write(Worker {
-                // Placeholder — overwritten by `init()` immediately below.
-                ctx: bun_ptr::BackRef::from(NonNull::<BundleV2<'static>>::dangling()),
+        let mut map = self.workers_assignments.lock();
+        let worker: *mut Worker = match map.entry(id) {
+            MapEntry::Occupied(o) => *o.into_mut(),
+            // Complete before it is inserted: `deinit_without_freeing_arena`
+            // takes this lock and tears down every Worker the map holds.
+            MapEntry::Vacant(v) => *v.insert(bun_core::heap::into_raw(Box::new(Worker {
+                // SAFETY: `v2` is the bundle that owns this pool; it outlives every Worker.
+                ctx: unsafe { bun_ptr::BackRef::from_raw(self.v2.cast_mut()) },
                 heap: None,
-                arena: bun_ptr::BackRef::from(NonNull::<ThreadLocalArena>::dangling()),
+                // Points into `heap` once `create()` runs.
+                arena: bun_ptr::BackRef::dangling(),
                 thread: NonNull::new(ThreadPoolLib::Thread::current())
                     .map(bun_ptr::ParentRef::from),
                 data: None,
@@ -446,11 +453,13 @@ impl ThreadPool {
                 },
                 temporary_arena: None,
                 stmt_list: None,
-            });
-            (*worker).init(&*self.v2);
-            TLS_WORKER.set((self.generation, worker));
-            &mut *worker
-        }
+            }))),
+        };
+        drop(map);
+        TLS_WORKER.set((self.generation, worker));
+        // SAFETY: the map only holds complete Workers; each lives until its
+        // `deinit_soon`, and only the thread whose id keys it uses it.
+        unsafe { &mut *worker }
     }
 }
 
@@ -664,12 +673,6 @@ impl Worker {
         self.ast_memory_store.pop();
     }
 
-    pub(crate) fn init(&mut self, v2: &BundleV2<'_>) {
-        // Lifetime-erase `'_` → `'static` via `NonNull::cast` (BACKREF: the
-        // bundle outlives every worker).
-        self.ctx = bun_ptr::BackRef::from(NonNull::from(v2).cast::<BundleV2<'static>>());
-    }
-
     fn create(&mut self, ctx: &BundleV2<'_>) {
         // `bun_perf::trace` takes a generated `PerfEvent` enum, and
         // the generator hasn't emitted `Bundler.Worker.create` yet (only
@@ -697,7 +700,6 @@ impl Worker {
         self.ast_memory_store.reset();
 
         let log: *mut bun_ast::Log = arena_ref.alloc(bun_ast::Log::init());
-        self.ctx = bun_ptr::BackRef::from(NonNull::from(ctx).cast::<BundleV2<'static>>());
         // Use a fresh Bump (no nested-arena type yet).
         self.temporary_arena = Some(bun_alloc::Arena::new());
         self.stmt_list = Some(StmtList::init());

src/bundler/bundle_v2.rs

@@ -2081,7 +2081,7 @@ pub mod bv2_impl {
             false
         }
 
-        pub(crate) fn wait_for_parse(&mut self) {
+        fn wait_for_parse(&mut self) {
             // `tick_raw` (not `tick`) — `is_done` reborrows `*ctx` as
             // `&mut BundleV2`, and `BundleV2` (via `linker.r#loop`) owns the
             // `AnyEventLoop` slot, so holding `&mut AnyEventLoop` across the
@@ -2108,6 +2108,18 @@ pub mod bv2_impl {
             );
         }
 
+        /// Runs one of the `enqueue_entry_points_*` variants and drains what it
+        /// scheduled, also when it failed: the runtime's parse task is on the pool
+        /// before anything in them can fail, and an error leads to the teardown.
+        fn enqueue_and_wait_for_parse(
+            &mut self,
+            enqueue: impl FnOnce(&mut Self) -> Result<(), Error>,
+        ) -> Result<(), Error> {
+            let enqueued = enqueue(self);
+            self.wait_for_parse();
+            enqueued
+        }
+
         /// Callers require an entry point, so none after parsing means one was dropped without an error.
         fn fail_if_no_entry_points(&self) -> Result<(), Error> {
             if !self.graph.entry_points.is_empty() {
@@ -3859,13 +3871,12 @@ pub mod bv2_impl {
 
                 let entry_points: *const [Box<[u8]>] =
                     &raw const *this.transpiler.options.entry_points;
-                // SAFETY: `transpiler.options.entry_points` is borrowed only for the duration
-                // of `enqueue_entry_points_normal`, which never frees/reallocates it; raw-ptr
-                // sidestep for the `&mut self` overlap.
-                this.enqueue_entry_points_normal(unsafe { &*entry_points })?;
-
-                // Like `run_from_js_in_new_thread`: drain the pool, then report entry point errors.
-                this.wait_for_parse();
+                this.enqueue_and_wait_for_parse(|bundle| {
+                    // SAFETY: `transpiler.options.entry_points` is borrowed only for the duration
+                    // of `enqueue_entry_points_normal`, which never frees/reallocates it; raw-ptr
+                    // sidestep for the `&mut self` overlap.
+                    bundle.enqueue_entry_points_normal(unsafe { &*entry_points })
+                })?;
                 this.dump_pool_stats("parse");
 
                 *minify_duration = (((bun_core::time::nano_timestamp() as i64)
@@ -3989,23 +4000,19 @@ pub mod bv2_impl {
             let mut this = BundleV2::init(transpiler, None, alloc, event_loop, false, None, alloc)?;
             this.unique_key = generate_unique_key();
 
-            if this.transpiler.log().has_errors() {
-                return Err(crate::Error::BuildFailed);
-            }
-
-            // enqueueEntryPoints schedules the runtime task before any fallible
-            // allocation. If a later allocation fails we must still drain the
-            // pool so workers aren't left holding pointers into the caller's
-            // stack-allocated Transpiler.
-            if let Err(err) = this.enqueue_entry_points_normal(entry_points) {
-                this.wait_for_parse();
+            // The caller tears down a returned bundle; an error exit has to do it here.
+            let scanned = if this.transpiler.log().has_errors() {
+                Err(crate::Error::BuildFailed)
+            } else {
+                this.enqueue_and_wait_for_parse(|bundle| {
+                    bundle.enqueue_entry_points_normal(entry_points)
+                })
+            };
+            if let Err(err) = scanned {
+                this.deinit_without_freeing_arena();
                 return Err(err);
             }
 
-            // Even if entry point resolution produced errors we still wait for
-            // all enqueued parse tasks to finish so the graph is consistent.
-            this.wait_for_parse();
-
             Ok(this)
         }
 
@@ -4034,10 +4041,9 @@ pub mod bv2_impl {
                     return Err(crate::Error::BuildFailed);
                 }
 
-                this.enqueue_entry_points_bake_production(entry_points)?;
-
-                // Drain the pool, then report entry point errors (as `generate_from_cli` does).
-                this.wait_for_parse();
+                this.enqueue_and_wait_for_parse(|bundle| {
+                    bundle.enqueue_entry_points_bake_production(entry_points)
+                })?;
 
                 if this.transpiler.log().has_errors() {
                     return Err(crate::Error::BuildFailed);
@@ -4882,6 +4888,19 @@ pub mod bv2_impl {
         }
 
         pub fn deinit_without_freeing_arena(&mut self) {
+            // Nothing may still run against the graph and the workers this frees.
+            // Parse results need the event loop, so the drivers drain those
+            // (`enqueue_and_wait_for_parse`; the dev server gets here from `is_done`).
+            assert_eq!(
+                self.graph.pending_items, 0,
+                "BundleV2 torn down with parse tasks in flight"
+            );
+            // Scheduled by `link`, waited for only by `generate_chunks_in_parallel`.
+            self.linker.source_maps.line_offset_wait_group.wait();
+            self.linker.source_maps.quoted_contents_wait_group.wait();
+            // The last IO callback may still be inside the worker pool `pool.deinit()` frees.
+            self.graph.pool().wait_for_io_tasks();
+
             {
                 // We do this first to make it harder for any dangling pointers to data to be used in there.
                 let on_parse_finalizers = core::mem::take(&mut self.finalizers);
@@ -5006,10 +5025,9 @@ pub mod bv2_impl {
 
             /* arena: help_catch_memory_issues — no-op (mimalloc TLH check) */
 
-            self.enqueue_entry_points_normal(entry_points)?;
-
-            // We must wait for all the parse tasks to complete, even if there are errors.
-            self.wait_for_parse();
+            self.enqueue_and_wait_for_parse(|bundle| {
+                bundle.enqueue_entry_points_normal(entry_points)
+            })?;
 
             /* arena: help_catch_memory_issues — no-op (mimalloc TLH check) */
 

src/bundler/linker_context/postProcessJSChunk.rs

@@ -104,7 +104,7 @@ pub(crate) fn post_process_js_chunk(
         ModuleInfo::create(loader.is_type_script())
     });
 
-    // SAFETY: worker.arena is set in Worker::init() before any task runs.
+    // `worker.arena` is set in `Worker::create()`, which `Worker::get` runs first.
     let worker_arena: &Arena = worker.arena();
 
     {

src/runtime/api/js_bundle_completion_task.rs

@@ -1248,23 +1248,16 @@ impl CompletionStruct for JSBundleCompletionTask {
             .map(|b| &**b)
             .collect();
 
-        let run = bv2.run_from_js_in_new_thread(&entry_points);
-
-        // The AST-allocator pop lives in `generate_in_new_thread`; the
-        // source-map wait-group waits run only on the error path.
-        match run {
+        // The AST-allocator pop lives in `generate_in_new_thread`.
+        let result = match bv2.run_from_js_in_new_thread(&entry_points) {
             Ok(build) => {
                 self.set_result(BundleV2Result::Value(build));
-                bv2.deinit_without_freeing_arena();
                 Ok(())
             }
-            Err(err) => {
-                bv2.linker.source_maps.line_offset_wait_group.wait();
-                bv2.linker.source_maps.quoted_contents_wait_group.wait();
-                bv2.deinit_without_freeing_arena();
-                Err(err)
-            }
-        }
+            Err(err) => Err(err),
+        };
+        bv2.deinit_without_freeing_arena();
+        result
     }
 }