Hand pool-finished transpiler and patch jobs back through their pointer, not a &mut receiver

src/install/patch_install.rs

@@ -141,24 +141,38 @@ impl PatchTask {
     /// # Safety
     /// Only invoked by `ThreadPool` via the `callback` fn-pointer registered in
     /// `new_calc_patch_hash` / `new_apply_patch_hash`. `task` must be live and
-    /// point at the `task` field of a heap-allocated `PatchTask`, with the pool
-    /// granting exclusive access for the duration of the call.
+    /// point at the `task` field of a heap-allocated `PatchTask`, which this
+    /// thread owns until it pushes the task back to the main thread.
     pub(crate) unsafe fn run_from_thread_pool(task: *mut ThreadPoolTask) {
         // SAFETY: thread-pool callback contract — `task` points to the `task`
         // field of a live `PatchTask` (set at construction); the pool runs
         // each task at most once with exclusive access for the call.
-        let patch_task = unsafe { &mut *PatchTask::from_task_ptr(task) };
-        patch_task.run_from_thread_pool_impl();
+        let this = unsafe { PatchTask::from_task_ptr(task) };
+        // SAFETY: as above. The main thread reclaims the task's `Box` as soon as
+        // the push below lands, so the `&mut` the autoref forms is scoped to
+        // its statement and the push goes through `this`, not a reference.
+        let mgr = unsafe {
+            (*this).run_from_thread_pool_impl();
+            (*this).manager.as_ptr()
+        };
+        // SAFETY: `this` is non-null (recovered from a live task). `mgr` is a
+        // long-lived BACKREF; the worker thread only touches the lock-free
+        // `patch_task_queue` and the event-loop wake atomics, neither of which
+        // alias data the main thread holds an exclusive borrow on.
+        unsafe {
+            (*mgr)
+                .patch_task_queue
+                .push(core::ptr::NonNull::new_unchecked(this));
+            PackageManager::wake_raw(mgr);
+        }
     }
 
-    pub(crate) fn run_from_thread_pool_impl(&mut self) {
+    fn run_from_thread_pool_impl(&mut self) {
         bun_output::scoped_log!(
             InstallPatch,
             "runFromThreadPoolImpl {}",
             <&'static str>::from(&self.callback)
         );
-        // There are no early returns in the body, so the ordering
-        // (body → push → wake) is inlined below.
         match &mut self.callback {
             Callback::CalcHash(_) => {
                 let result = self.calc_hash();
@@ -173,17 +187,6 @@ impl PatchTask {
                 self.apply().expect("OOM");
             }
         }
-        let mgr = self.manager.as_ptr();
-        // SAFETY: `self.manager` is a long-lived BACKREF;
-        // the worker thread only touches the lock-free `patch_task_queue` and the
-        // event-loop wake atomics, neither of which alias data the main thread
-        // holds an exclusive borrow on.
-        unsafe {
-            (*mgr)
-                .patch_task_queue
-                .push(core::ptr::NonNull::from(&mut *self));
-            PackageManager::wake_raw(mgr);
-        }
     }
 
     pub(crate) fn run_from_main_thread(

src/jsc/RuntimeTranspilerStore.rs

@@ -272,8 +272,9 @@ impl RuntimeTranspilerStore {
             return;
         }
         // we run just one job first to see if there are more
-        // SAFETY: `first` is a live job popped from the intrusive queue.
-        if let Err(err) = unsafe { (*first).run_from_js_thread() } {
+        // SAFETY: `first` is a live job popped from the intrusive queue; the
+        // batch iterator has already moved past it.
+        if let Err(err) = unsafe { TranspilerJob::run_from_js_thread(first) } {
             global.report_uncaught_exception_from_error(err);
         }
         loop {
@@ -288,8 +289,8 @@ impl RuntimeTranspilerStore {
             {
                 return;
             }
-            // SAFETY: `job` is a live job popped from the intrusive queue.
-            if let Err(err) = unsafe { (*job).run_from_js_thread() } {
+            // SAFETY: as for `first`.
+            if let Err(err) = unsafe { TranspilerJob::run_from_js_thread(job) } {
                 global.report_uncaught_exception_from_error(err);
             }
         }
@@ -403,9 +404,10 @@ pub struct TranspilerJob {
     // raw pointers/BackRefs are used (BACKREF — VM owns the
     // store and outlives every job).
     pub(crate) vm: *mut VirtualMachine,
-    /// The pool thread runs this job under `loop_handle.borrow()`: the job's
-    /// own slot, the transpiler it copies and the store queue it pushes to are
-    /// all VM-owned, and the VM's teardown waits for the borrow to end.
+    /// The pool thread transpiles under `loop_handle.borrow_if_running()` (the
+    /// transpiler it copies is VM-owned) and counts as embedded work from
+    /// `schedule` until it has pushed the slot back to the store queue, so the
+    /// VM's teardown waits for both.
     pub(crate) loop_handle: crate::LoopHandle,
     pub global_this: BackRef<JSGlobalObject>,
     pub(crate) fetcher: Fetcher,
@@ -447,6 +449,19 @@ impl Fetcher {
     }
 }
 
+/// A finished job's result, moved out of its slot (`TranspilerJob::take_completion`)
+/// so the slot can be returned to the store before `AsyncModule::fulfill` runs.
+/// `specifier` and `referrer` carry the +1 that `fulfill` releases.
+struct Completion {
+    promise: JSValue,
+    global_this: BackRef<JSGlobalObject>,
+    specifier: String,
+    referrer: String,
+    log: bun_ast::Log,
+    resolved_source: OwnedResolvedSource,
+    parse_error: Option<crate::CrateError>,
+}
+
 /// Per-worker output buffer. The printer is the **only** state
 /// retained across `run()` calls — its backing `Vec<u8>` is genuinely worth
 /// reusing (capped at 512 K / 2 M below). The parse arena and AST memory
@@ -484,8 +499,9 @@ fn tls_get_or_leak<T>(
 
 impl TranspilerJob {
     /// Kept as a private inherent fn (not `impl Drop`) because the
-    /// slot is recycled into the HiveArray via `store.put(this)`. Only caller is
-    /// `run_from_js_thread`.
+    /// slot is recycled into the HiveArray via `store.put(this)`. Called right
+    /// before that `put` by `take_completion` and
+    /// `release_queued_jobs_for_teardown`.
     ///
     /// Note: `HiveArrayFallback::put` runs `drop_in_place` on the slot (see
     /// hive_array.rs note), so the Drop-carrying fields — `OwnedString` ×2,
@@ -514,40 +530,86 @@ impl TranspilerJob {
         // replacement a second time).
     }
 
-    fn dispatch_to_main_thread(&mut self) {
-        let vm = self.vm;
-        let loop_handle = self.loop_handle.clone();
+    /// Hands the job back to the JS thread, which completes it
+    /// (`run_from_js_thread`) or releases it (`release_queued_jobs_for_teardown`);
+    /// either way it writes the slot and `put()`s it (a `Box` free once the hive
+    /// is full) as soon as the push lands. Takes the pointer, not `&mut self`: a
+    /// reference argument stays protected until its call returns, so the JS
+    /// thread would be writing to, or freeing, protected memory. The caller must
+    /// not hold a reference into `*this` across the call either.
+    ///
+    /// # Safety
+    /// `this` is the live slot `transpile()` scheduled, still owned by the pool
+    /// thread; nothing on this thread touches `*this` afterwards.
+    unsafe fn dispatch_to_main_thread(this: *mut Self) {
+        // SAFETY: fn contract; both accesses end at the `;`, before the push.
+        let (vm, loop_handle) = unsafe { ((*this).vm, (*this).loop_handle.clone()) };
         // SAFETY: vm outlives the job (BACKREF — VM owns the store).
         let transpiler_store: *mut RuntimeTranspilerStore =
             unsafe { ptr::addr_of_mut!((*vm).transpiler_store) };
-        let job = NonNull::from(&mut *self);
-        // SAFETY: queue is concurrent-safe (UnboundedQueue uses atomics).
-        unsafe { (*transpiler_store).queue.push(job) };
-        // Another thread may free `self` at any time after .push, so we cannot use it any more
-        // (the handle was cloned out above for exactly this reason). The VM
-        // waits for embedded work before closing its handle, so this is queued.
+        // SAFETY: `this` is non-null per fn contract; the queue is
+        // concurrent-safe (UnboundedQueue uses atomics). The JS thread owns the
+        // job from here on.
+        unsafe { (*transpiler_store).queue.push(NonNull::new_unchecked(this)) };
+        // The VM waits for embedded work before closing its handle, so this is
+        // queued.
         let crate::vm_handle::Posted::Queued =
             loop_handle.post_task(ConcurrentTask::create_from(transpiler_store))
         else {
             unreachable!("VM handle closed with embedded transpile work outstanding");
         };
     }
 
-    fn run_from_js_thread(&mut self) -> JsResult<()> {
-        let vm = self.vm;
+    /// Completes the import. The slot goes back to the store first (a `Box`
+    /// free when it was heap-spilled; `fulfill` runs script, which may claim it
+    /// again), so this takes the pointer rather than a `&mut self` that would
+    /// still be protected while the slot is dropped or freed.
+    ///
+    /// # Safety
+    /// `this` is a live job popped from the store queue; the caller does not
+    /// touch it afterwards.
+    unsafe fn run_from_js_thread(this: *mut Self) -> JsResult<()> {
+        // SAFETY: fn contract; both accesses end at the `;`, before the put.
+        let (vm, completion) = unsafe { ((*this).vm, (*this).take_completion()) };
+        // SAFETY: `vm` outlives the job; `this` is the slot `transpile()` claimed
+        // from this store and nothing uses it past this point.
+        unsafe { (*vm).transpiler_store.store.put(this) };
+
+        let Completion {
+            promise,
+            global_this,
+            specifier,
+            referrer,
+            mut log,
+            resolved_source,
+            parse_error,
+        } = completion;
+        let mut resolved_source = resolved_source.into_ffi();
+        AsyncModule::fulfill(
+            &global_this,
+            promise,
+            &mut resolved_source,
+            parse_error,
+            specifier,
+            referrer,
+            &mut log,
+        )
+    }
+
+    /// Moves everything `AsyncModule::fulfill` needs out of the slot and
+    /// resets it, leaving only what `store.put()`'s drop glue releases.
+    fn take_completion(&mut self) -> Completion {
         let promise = self.promise.swap();
-        // Copy the BackRef out (it is `Copy`) so the borrow of `*self` ends
-        // before `reset_for_pool`/`put` need `&mut *self` below; deref at the
-        // `fulfill` call site instead.
         let global_this = self.global_this;
         // Note: the KeepAlive takes an `EventLoopCtx`
         // vtable; resolve it via the `get_vm_ctx` hook (registered by `bun_runtime::init`).
         self.poll_ref.unref(get_vm_ctx(AllocatorType::Js));
 
         let referrer = core::mem::take(&mut self.non_threadsafe_referrer).into_inner();
-        let mut log = core::mem::replace(&mut self.log, bun_ast::Log::init());
-        // Take RAII ownership out of the job; `into_ffi()` below transfers the
-        // +1 strings to `AsyncModule::fulfill` → C++ `Zig::ResolvedSource`.
+        let log = core::mem::replace(&mut self.log, bun_ast::Log::init());
+        // Take RAII ownership out of the job; `into_ffi()` in the caller
+        // transfers the +1 strings to `AsyncModule::fulfill` → C++
+        // `Zig::ResolvedSource`.
         let mut owned_resolved_source = core::mem::take(&mut self.resolved_source);
         let resolved_source = owned_resolved_source.as_mut();
         let specifier = 'brk: {
@@ -569,24 +631,15 @@ impl TranspilerJob {
         self.promise.deinit();
         self.reset_for_pool();
 
-        // SAFETY: vm outlives the job; transpiler_store.store.put recycles the slot.
-        unsafe {
-            (*vm)
-                .transpiler_store
-                .store
-                .put(std::ptr::from_mut::<TranspilerJob>(self))
-        };
-
-        let mut resolved_source = owned_resolved_source.into_ffi();
-        AsyncModule::fulfill(
-            &global_this,
+        Completion {
             promise,
-            &mut resolved_source,
-            parse_error,
+            global_this,
             specifier,
             referrer,
-            &mut log,
-        )
+            log,
+            resolved_source: owned_resolved_source,
+            parse_error,
+        }
     }
 
     fn schedule(&mut self) {
@@ -613,16 +666,18 @@ impl TranspilerJob {
         // SAFETY: as above.
         let handle = unsafe { (*this).loop_handle.clone() };
         if let Some(_vm) = handle.borrow_if_running() {
-            // SAFETY: live slot, exclusively ours until dispatched.
+            // SAFETY: live slot, exclusively ours until dispatched. The `&mut`
+            // this autoref forms ends with the statement, before the dispatch
+            // publishes the slot.
             unsafe { (*this).run() };
-        } else {
-            // SAFETY: as above.
-            unsafe { (*this).dispatch_to_main_thread() };
         }
-        // Last touch of the slot from this thread was the dispatch.
+        // SAFETY: as above; the dispatch is this thread's last touch of the slot.
+        unsafe { Self::dispatch_to_main_thread(this) };
         handle.embedded_work_finished();
     }
 
+    /// Fills in `resolved_source` or `parse_error`; the caller hands the job
+    /// back to the JS thread afterwards, on every return path.
     fn run(&mut self) {
         // Stack-local per call, bulk-freed on return. An earlier version hoisted
         // this to a per-worker-thread leaked `Box<MimallocArena>` (and a second
@@ -640,14 +695,6 @@ impl TranspilerJob {
         // between calls.
         let arena = Arena::new();
 
-        // `defer this.dispatchToMainThread()` — fires on every return path.
-        let this_ptr: *mut TranspilerJob = self;
-        scopeguard::defer! {
-            // SAFETY: `self` outlives this guard (guard drops before fn return);
-            // no other &mut alias is live at drop time.
-            unsafe { (*this_ptr).dispatch_to_main_thread() };
-        }
-
         // SAFETY contract: `vm` outlives the job (BACKREF — VM owns the store).
         // Note: kept as a raw pointer — never form `&mut VirtualMachine`
         // here. (a) the JS thread is concurrently live on the same VM, so a