napi: drive napi_async_work through the addon's pointer instead of &mut self receivers

src/runtime/dispatch.rs

@@ -370,7 +370,10 @@ pub(crate) fn run_task(
 
         // ── napi ─────────────────────────────────────────────────────────
         task_tag::NapiAsyncWork => {
-            cast!(napi_async_work).run_from_js(vm, global);
+            // SAFETY: §Dispatch — tag identifies the pointee; the addon's
+            // `complete` callback usually frees the work, so it takes the raw
+            // pointer (no `&mut` at this boundary).
+            unsafe { napi_async_work::run_from_js(cast_ptr!(napi_async_work), vm, global) };
         }
         task_tag::ThreadSafeFunction => {
             ThreadSafeFunction::on_dispatch(cast_ptr!(ThreadSafeFunction));

src/runtime/napi/napi_body.rs

@@ -50,7 +50,7 @@ impl Taskable for napi_async_work {
         let vm = VirtualMachine::get().as_mut();
         let global = vm.global();
         // SAFETY: fn contract — the addon's live work object the pool posted.
-        unsafe { (*this).run_from_js(vm, global) };
+        unsafe { napi_async_work::run_from_js(this, vm, global) };
     }
 }
 impl Taskable for ThreadSafeFunction {
@@ -1811,111 +1811,169 @@ impl napi_async_work {
         drop(unsafe { bun_core::heap::take(this) });
     }
 
-    pub(crate) fn schedule(&mut self) {
-        if self.scheduled {
-            return;
+    // The work is shared between the addon (which owns and frees it), the
+    // pool thread and the event-loop queue, so every entry point below takes
+    // the addon's pointer and touches individual fields through it: no frame
+    // holds a reference to the whole work while another party may write to
+    // it (`cancel` racing `run`) or free it (`complete`, which normally calls
+    // `napi_delete_async_work`, runs as soon as `run` has posted).
+
+    /// `napi_queue_async_work`, JS thread.
+    ///
+    /// # Safety
+    /// `this` is a live work from [`Self::new`].
+    pub(crate) unsafe fn schedule(this: *mut Self) {
+        // SAFETY: fn contract; the pool only gets the work at the end of this
+        // function, so until then the JS thread is the only one touching it.
+        unsafe {
+            if (*this).scheduled {
+                return;
+            }
+            (*this).scheduled = true;
+            (*this).poll_ref.ref_(bun_io::js_vm_ctx());
+            // The work object belongs to the addon and `execute` receives this
+            // env: counted, so the VM waits for it (Node likewise settles its
+            // threadpool requests before an environment is freed).
+            (*this).loop_handle.embedded_work_scheduled();
+            // Projected from the work pointer itself, as `from_task_ptr`
+            // requires of the pointer it gets back.
+            WorkPool::schedule(&raw mut (*this).task);
         }
-        self.scheduled = true;
-        self.poll_ref.ref_(bun_io::js_vm_ctx());
-        // The work object belongs to the addon and `execute` receives this
-        // env: counted, so the VM waits for it (Node likewise settles its
-        // threadpool requests before an environment is freed).
-        self.loop_handle.embedded_work_scheduled();
-        WorkPool::schedule(&raw mut self.task);
     }
 
     pub(crate) unsafe fn run_from_thread_pool(task: *mut WorkPoolTask) {
-        // SAFETY: `task` is the `task` field of a live heap `napi_async_work`,
-        // exclusively owned by the work pool for this callback's duration.
-        unsafe { (*napi_async_work::from_task_ptr(task)).run() };
+        // SAFETY: `task` is the field `schedule` handed to the pool, projected
+        // from the work pointer, and the pool runs each task once.
+        unsafe { Self::run(napi_async_work::from_task_ptr(task)) };
     }
 
-    fn run(&mut self) {
-        let self_ptr: *mut Self = self;
-        let handle = self.loop_handle.clone();
+    /// Pool thread. Ends by handing the work to the JS thread, which runs
+    /// `complete` (and with it, usually, `napi_delete_async_work`) as soon as
+    /// the post lands, so nothing reads `*this` after the post; the cloned
+    /// `handle` is what the pool thread keeps. `napi_cancel_async_work` may
+    /// write `status` at any point during this function.
+    ///
+    /// # Safety
+    /// `this` is the live work `schedule` handed to the pool; called once per
+    /// scheduling.
+    unsafe fn run(this: *mut Self) {
+        // SAFETY: fn contract. The JS thread does not touch these fields while
+        // the work is in the pool's hands (`status` only through its atomic).
+        let handle = unsafe { (*this).loop_handle.clone() };
         // A VM that is already stopping cancels work it has not started, as
         // Node's environment cleanup does (uv_cancel); otherwise `execute` runs
         // with the VM held open.
         let vm = handle.borrow_if_running();
         let started = vm.is_some()
-            && match self.status.compare_exchange(
-                AsyncWorkStatus::Pending as u32,
-                AsyncWorkStatus::Started as u32,
-                Ordering::SeqCst,
-                Ordering::SeqCst,
-            ) {
+            // SAFETY: as above.
+            && match unsafe {
+                (*this).status.compare_exchange(
+                    AsyncWorkStatus::Pending as u32,
+                    AsyncWorkStatus::Started as u32,
+                    Ordering::SeqCst,
+                    Ordering::SeqCst,
+                )
+            } {
                 Ok(_) => true,
                 Err(state) => state != AsyncWorkStatus::Cancelled as u32,
             };
         if started {
-            (self.execute)(self.env.get(), self.data);
-            self.status
-                .store(AsyncWorkStatus::Completed as u32, Ordering::SeqCst);
+            // SAFETY: as above.
+            let (execute, env, data) =
+                unsafe { ((*this).execute, (*this).env.get(), (*this).data) };
+            execute(env, data);
+            // SAFETY: as above.
+            unsafe {
+                (*this)
+                    .status
+                    .store(AsyncWorkStatus::Completed as u32, Ordering::SeqCst)
+            };
         } else {
-            let _ = self.cancel();
+            // SAFETY: as above.
+            let _ = unsafe { Self::cancel(this) };
         }
         drop(vm);
-        self.post_to_js_thread(self_ptr);
-        // `self` may already be freed by the JS thread; the handle is ours.
-        handle.embedded_work_finished();
-    }
-
-    /// Pool thread → JS thread: run `complete` there. `concurrent_task` is the
-    /// live inline field of this heap work; the queue takes ownership of its
-    /// `next` link. Counted work, so the VM has not closed its handle; a VM
-    /// tearing down runs `complete` from its queue release (status cancelled
-    /// if `execute` never ran), as Node does at environment cleanup.
-    fn post_to_js_thread(&mut self, self_ptr: *mut Self) {
-        let ct = core::ptr::NonNull::from(
-            self.concurrent_task
-                .from(self_ptr, AutoDeinit::ManualDeinit),
-        );
-        let bun_jsc::vm_handle::Posted::Queued = self.loop_handle.post_task(ct) else {
+        // `concurrent_task` is the inline field of this heap work; the queue
+        // takes ownership of its `next` link. Counted work, so the VM has not
+        // closed its handle; a VM tearing down runs `complete` from its queue
+        // release (status cancelled if `execute` never ran), as Node does at
+        // environment cleanup.
+        // SAFETY: as above; this is the pool thread's last access to `*this`.
+        let ct = core::ptr::NonNull::from(unsafe {
+            (*this).concurrent_task.from(this, AutoDeinit::ManualDeinit)
+        });
+        let bun_jsc::vm_handle::Posted::Queued = handle.post_task(ct) else {
             unreachable!("VM handle closed with napi async work outstanding");
         };
+        // The JS thread may have freed the work by now; `handle` is ours.
+        handle.embedded_work_finished();
     }
 
-    pub(crate) fn cancel(&mut self) -> bool {
-        self.status
-            .compare_exchange(
+    /// `napi_cancel_async_work` (JS thread) and [`Self::run`] (pool thread),
+    /// possibly at the same time: only the atomic is touched.
+    ///
+    /// # Safety
+    /// `this` is a live work.
+    pub(crate) unsafe fn cancel(this: *mut Self) -> bool {
+        // SAFETY: fn contract.
+        unsafe {
+            (*this).status.compare_exchange(
                 AsyncWorkStatus::Pending as u32,
                 AsyncWorkStatus::Cancelled as u32,
                 Ordering::SeqCst,
                 Ordering::SeqCst,
             )
-            .is_ok()
+        }
+        .is_ok()
     }
 
-    pub(crate) fn run_from_js(&mut self, vm: &mut VirtualMachine, global: &JSGlobalObject) {
-        // Note: the "this" value here may already be freed by the user in `complete`
-        // Note: KeepAlive is not `Copy`, so move it out (the original slot may
-        // be freed under us by `complete`).
-        let mut poll_ref = core::mem::take(&mut self.poll_ref);
+    /// JS thread, once per scheduling: from the task queue, or from its
+    /// release when the VM tears down first. `complete` is the addon's and
+    /// normally calls `napi_delete_async_work` on this very work, so
+    /// everything the call needs is taken out of `*this` first and `this` is
+    /// not used again after `complete` starts.
+    ///
+    /// # Safety
+    /// `this` is the live work [`Self::run`] posted; the pool thread is done
+    /// with it.
+    pub(crate) unsafe fn run_from_js(
+        this: *mut Self,
+        vm: &mut VirtualMachine,
+        global: &JSGlobalObject,
+    ) {
+        // SAFETY: fn contract; the JS thread is the only one touching the work now.
+        let (mut poll_ref, complete, env, status, data) = unsafe {
+            (
+                core::mem::take(&mut (*this).poll_ref),
+                (*this).complete,
+                (*this).env.get(),
+                (*this).status.load(Ordering::SeqCst),
+                (*this).data,
+            )
+        };
         // KeepAlive::unref needs an event-loop ctx so it cannot impl Drop
         // generically; this is a genuine one-off cleanup.
         scopeguard::defer! { poll_ref.unref(bun_io::js_vm_ctx()); }
 
         // https://github.com/nodejs/node/blob/a2de5b9150da60c77144bb5333371eaca3fab936/src/node_api.cc#L1201
-        let Some(complete) = self.complete else {
+        let Some(complete) = complete else {
             return;
         };
 
-        let env = self.env.get();
-        // SAFETY: env is held alive by NapiEnvRef for the duration of this call.
+        // SAFETY: `global` (live for this call) holds a ref on every env made
+        // for it (`GlobalObject::m_napiEnvs`), so `env` outlives the work's own
+        // ref, which `complete` usually drops by freeing the work.
         let env_ref = unsafe { &*env };
         let _hs = NapiHandleScope::open_scoped(env_ref);
 
-        let status: NapiStatus =
-            if self.status.load(Ordering::SeqCst) == AsyncWorkStatus::Cancelled as u32 {
-                NapiStatus::cancelled
-            } else {
-                NapiStatus::ok
-            };
+        let status: NapiStatus = if status == AsyncWorkStatus::Cancelled as u32 {
+            NapiStatus::cancelled
+        } else {
+            NapiStatus::ok
+        };
 
-        complete(env, status as napi_status, self.data);
+        complete(env, status as napi_status, data);
 
-        // SAFETY: env is valid for the duration of this call.
-        let env_ref = unsafe { &*env };
         if let Some(exception) = env_ref.get_and_clear_pending_exception() {
             let _ = vm.uncaught_exception(global, exception, false);
         } else if global.has_exception() {
@@ -2172,11 +2230,13 @@ extern "C" fn napi_create_async_work(
 extern "C" fn napi_delete_async_work(env_: napi_env, work_: *mut napi_async_work) -> napi_status {
     bun_output::scoped_log!(napi, "napi_delete_async_work");
     let env = get_env!(env_);
-    // SAFETY: `work_` is null or the `napi_async_work` we allocated in `napi_create_async_work`.
-    let Some(work) = (unsafe { work_.as_mut() }) else {
+    if work_.is_null() {
         return env.invalid_arg();
-    };
-    debug_assert!(core::ptr::eq(env.to_js(), work.global.as_ptr()));
+    }
+    // SAFETY: non-null `work_` is the work `napi_create_async_work` allocated.
+    debug_assert!(core::ptr::eq(env.to_js(), unsafe {
+        (*work_).global.as_ptr()
+    }));
     napi_async_work::destroy(work_);
     env.ok()
 }
@@ -2185,25 +2245,33 @@ extern "C" fn napi_delete_async_work(env_: napi_env, work_: *mut napi_async_work
 extern "C" fn napi_queue_async_work(env_: napi_env, work_: *mut napi_async_work) -> napi_status {
     bun_output::scoped_log!(napi, "napi_queue_async_work");
     let env = get_env!(env_);
-    // SAFETY: `work_` is null or the `napi_async_work` we allocated in `napi_create_async_work`.
-    let Some(work) = (unsafe { work_.as_mut() }) else {
+    if work_.is_null() {
         return env.invalid_arg();
-    };
-    debug_assert!(core::ptr::eq(env.to_js(), work.global.as_ptr()));
-    work.schedule();
+    }
+    // SAFETY: non-null `work_` is the work `napi_create_async_work` allocated.
+    debug_assert!(core::ptr::eq(env.to_js(), unsafe {
+        (*work_).global.as_ptr()
+    }));
+    // SAFETY: as above.
+    unsafe { napi_async_work::schedule(work_) };
     env.ok()
 }
 
 #[unsafe(no_mangle)]
 extern "C" fn napi_cancel_async_work(env_: napi_env, work_: *mut napi_async_work) -> napi_status {
     bun_output::scoped_log!(napi, "napi_cancel_async_work");
     let env = get_env!(env_);
-    // SAFETY: `work_` is null or the `napi_async_work` we allocated in `napi_create_async_work`.
-    let Some(work) = (unsafe { work_.as_mut() }) else {
+    if work_.is_null() {
         return env.invalid_arg();
-    };
-    debug_assert!(core::ptr::eq(env.to_js(), work.global.as_ptr()));
-    if work.cancel() {
+    }
+    // SAFETY: non-null `work_` is the work `napi_create_async_work` allocated;
+    // the pool thread may be running it, which is why nothing here forms a
+    // reference to the whole work.
+    debug_assert!(core::ptr::eq(env.to_js(), unsafe {
+        (*work_).global.as_ptr()
+    }));
+    // SAFETY: as above.
+    if unsafe { napi_async_work::cancel(work_) } {
         return env.ok();
     }