napi: keep dispatching threadsafe function calls when a callback blocks in a nested event loop

src/runtime/napi/napi_body.rs

@@ -2410,6 +2410,14 @@
 
     pub callback: TsfnCallback,
     pub(crate) dispatch_state: AtomicU8, // DispatchState
+    /// JS-thread only: depth of live (possibly nested) `on_dispatch` frames.
+    pub(crate) dispatch_depth: u32,
+    /// JS-thread only: `Closed` was observed while frames or tasks still
+    /// reference this object; whoever is last out frees it.
+    pub(crate) pending_destroy: bool,
+    /// Queued event-loop tasks targeting `on_dispatch`; atomic because addon
+    /// threads schedule them.
+    pub(crate) inflight_dispatch_tasks: AtomicU32,
     pub(crate) blocking_condvar: Condvar,
     pub(crate) closing: AtomicU8, // ClosingState
     /// Written under `lock` by `env_teardown` on the JS thread. Every path
@@ -2506,21 +2514,39 @@
     #[allow(clippy::not_unsafe_ptr_arg_deref)]
     pub(crate) fn on_dispatch(this: *mut ThreadSafeFunction) {
         // SAFETY: `this` is a live heap allocation owned by the event loop
-        // dispatch; `env_dead` is atomic so a shared reborrow suffices.
+        // dispatch; destroy is deferred while `dispatch_depth > 0` or tasks
+        // remain in flight. No `&mut *this` is held across `dispatch_one`,
+        // which runs user JS that can re-enter this dispatch.
+        let inflight = unsafe {
+            (*this)
+                .inflight_dispatch_tasks
+                .fetch_sub(1, Ordering::SeqCst)
+        } - 1;
+        // SAFETY: as above.
         if unsafe { (*this).env_dead.load(Ordering::SeqCst) } {
             // `env_teardown` already released everything and owns the free
             // decision. The loop this task came from is being destroyed.
             return;
         }
         // SAFETY: as above.
         if unsafe { (*this).closing.load(Ordering::SeqCst) } == ClosingState::Closed as u8 {
+            // SAFETY: as above.
+            if unsafe { (*this).dispatch_depth } > 0 || inflight > 0 {
+                // An outer frame or queued task still references `*this`;
+                // whoever is last out frees it.
+                // SAFETY: as above.
+                unsafe { (*this).pending_destroy = true };
+                return;
+            }
             // Finalize the ThreadSafeFunction.
             // SAFETY: `this` is the live heap allocation we own; closed state guarantees no other thread will touch it.
             unsafe { ThreadSafeFunction::destroy(this) };
             return;
         }
 
         let mut is_first = true;
+        // SAFETY: as above.
+        unsafe { (*this).dispatch_depth += 1 };
 
         // Run the tasks.
         loop {
@@ -2530,9 +2556,8 @@
                     .dispatch_state
                     .store(DispatchState::Running as u8, Ordering::SeqCst)
             };
-            // SAFETY: as above. `dispatch_one` runs JS that can re-enter other
-            // TSFN entry points, so the exclusive borrow is scoped to this call.
-            if unsafe { (*this).dispatch_one(is_first) } {
+            // SAFETY: as above.
+            if unsafe { Self::dispatch_one(this, is_first) } {
                 is_first = false;
                 // SAFETY: as above.
                 unsafe {
@@ -2541,16 +2566,8 @@
                         .store(DispatchState::Pending as u8, Ordering::SeqCst)
                 };
             } else {
-                // We're done running tasks, for now. Transition Running → Idle
-                // via CAS instead of an unconditional store: between
-                // dispatch_one() observing an empty queue (and dropping the
-                // lock) and this point, another thread may have enqueued an
-                // item and called schedule_dispatch(). That swap() saw
-                // Running, so it intentionally did *not* schedule a new
-                // concurrent task — it relies on this loop to pick the item
-                // up. If we blindly stored Idle we'd overwrite that Pending
-                // and the callback would be dropped (flaky lost-wakeup under
-                // load). On CAS failure, loop and re-drain.
+                // CAS, not a store: a concurrent schedule_dispatch() may have
+                // set Pending, and storing Idle over it would drop a wakeup.
                 // SAFETY: as above.
                 if unsafe {
                     (*this).dispatch_state.compare_exchange(
@@ -2568,6 +2585,19 @@
             }
         }
 
+        // SAFETY: as above.
+        let destroy_now = unsafe {
+            (*this).dispatch_depth -= 1;
+            (*this).dispatch_depth == 0
+                && (*this).pending_destroy
+                && (*this).inflight_dispatch_tasks.load(Ordering::SeqCst) == 0
+        };
+        if destroy_now {
+            // SAFETY: outermost frame, no other task references `*this`.
+            unsafe { ThreadSafeFunction::destroy(this) };
+            return;
+        }
+
         // Node sets a maximum number of runs per ThreadSafeFunction to 1,000.
         // We don't set a max. I would like to see an issue caused by not
         // setting a max before we do set a max. It is better for performance to
@@ -2603,8 +2633,11 @@
                     self.callback = TsfnCallback::Js(StrongOptional::empty());
                     self.poll_ref.disable();
                     let self_ptr: *mut Self = self;
+                    // The finalize task targets `on_dispatch` too.
+                    let _ = self.inflight_dispatch_tasks.fetch_add(1, Ordering::SeqCst);
                     let Some(loop_) = self.loop_mut() else {
                         // env torn down: `env_teardown` owns the finalize + free.
+                        let _ = self.inflight_dispatch_tasks.fetch_sub(1, Ordering::SeqCst);
                         return;
                     };
                     loop_.enqueue_task(Task::init(self_ptr));
@@ -2617,101 +2650,158 @@
         }
     }
 
-    pub(crate) fn dispatch_one(&mut self, is_first: bool) -> bool {
-        let mut queue_finalizer_after_call = false;
-        let task = 'brk: {
-            // `MutexGuard` holds the lock by raw pointer, so it does not borrow
-            // `*self` across the `&mut self` calls below.
-            let _g = self.lock.lock_guard();
-            let was_blocked = self.queue.is_blocked();
-            let Some(t) = self.queue.data.read_item() else {
-                // When there are no tasks and the number of threads that have
-                // references reaches zero, we prepare to finalize the
-                // ThreadSafeFunction.
-                if self.thread_count.load(Ordering::SeqCst) == 0 {
-                    if self.queue.max_queue_size > 0 {
-                        self.blocking_condvar.signal();
-                    }
-                    self.maybe_queue_finalizer();
-                }
-                return false;
-            };
-
-            if self.queue.count.fetch_sub(1, Ordering::SeqCst) == 1
-                && self.thread_count.load(Ordering::SeqCst) == 0
-            {
-                self.closing
-                    .store(ClosingState::Closing as u8, Ordering::SeqCst);
+    /// Dequeues one item under the lock, running the empty-queue finalize
+    /// check and scheduling the backup dispatch as needed. Never enters JS.
+    fn take_one_locked(&mut self, queue_finalizer_after_call: &mut bool) -> Option<*mut c_void> {
+        let _g = self.lock.lock_guard();
+        let was_blocked = self.queue.is_blocked();
+        let Some(t) = self.queue.data.read_item() else {
+            // When there are no tasks and the number of threads that have
+            // references reaches zero, we prepare to finalize the
+            // ThreadSafeFunction.
+            if self.thread_count.load(Ordering::SeqCst) == 0 {
                 if self.queue.max_queue_size > 0 {
                     self.blocking_condvar.signal();
                 }
-                queue_finalizer_after_call = true;
-            } else if was_blocked && !self.queue.is_blocked() {
+                self.maybe_queue_finalizer();
+            }
+            return None;
+        };
+
+        let prev_count = self.queue.count.fetch_sub(1, Ordering::SeqCst);
+        if prev_count == 1 && self.thread_count.load(Ordering::SeqCst) == 0 {
+            self.closing
+                .store(ClosingState::Closing as u8, Ordering::SeqCst);
+            if self.queue.max_queue_size > 0 {
                 self.blocking_condvar.signal();
             }
+            *queue_finalizer_after_call = true;
+        } else if was_blocked && !self.queue.is_blocked() {
+            self.blocking_condvar.signal();
+        }
 
-            break 'brk t;
-        };
+        if prev_count > 1 && self.inflight_dispatch_tasks.load(Ordering::SeqCst) == 0 {
+            // `call` can block in a nested event loop (#36828); one backup
+            // dispatch keeps the queued-behind items reachable.
+            self.schedule_dispatch();
+        }
 
-        if self.call(task, is_first).is_err() {
-            return false;
+        Some(t)
+    }
+
+    /// Runs one queued call. `this` stays raw: the microtask drain and `call`
+    /// enter user JS, which can re-enter this function's dispatch through a
+    /// nested event loop.
+    ///
+    /// # Safety
+    /// `this` is a live threadsafe function on the JS thread.
+    unsafe fn dispatch_one(this: *mut Self, is_first: bool) -> bool {
+        // Drain microtasks between callbacks (node#38506) BEFORE dequeuing:
+        // if the drain blocks in a nested event loop, the next item is still
+        // in the queue for a nested dispatch to deliver, in push order.
+        if !is_first {
+            let loop_: *mut EventLoop = {
+                // SAFETY: scoped reborrow.
+                match unsafe { (*this).event_loop.as_mut() } {
+                    // SAFETY: BackRef invariant while `Some`; JS thread,
+                    // outside tick().
+                    Some(back_ref) => unsafe { back_ref.get_mut() },
+                    None => return false,
+                }
+            };
+            // SAFETY: the per-thread event loop outlives this call.
+            if unsafe { (*loop_).drain_microtasks() }.is_err() {
+                return false;
+            }
         }
 
+        let mut queue_finalizer_after_call = false;
+        // SAFETY: call-scoped reborrow; `take_one_locked` never enters JS.
+        let Some(task) = (unsafe { (*this).take_one_locked(&mut queue_finalizer_after_call) })
+        else {
+            return false;
+        };
+
+        // No borrow of `*this` is live while user JS runs.
+        // SAFETY: per fn contract.
+        unsafe { Self::call(this, task) };
+
         if queue_finalizer_after_call {
-            self.maybe_queue_finalizer();
+            // SAFETY: call-scoped reborrow; `this` is still live (destroy deferred).
+            unsafe { (*this).maybe_queue_finalizer() };
         }
 
         // An item was dequeued: keep on_dispatch looping so remaining queued
         // items drain and the empty-queue thread_count==0 path can finalize.
         true
     }
 
-    /// This function can be called multiple times in one tick of the event loop.
-    /// See: https://github.com/nodejs/node/pull/38506
-    /// In that case, we need to drain microtasks.
-    fn call(&mut self, task: *mut c_void, is_first: bool) -> Result<(), bun_jsc::JsTerminated> {
-        let Some(env) = self.env.as_ref().map(NapiEnvRef::get) else {
+    /// Invokes the callback for one item. `this` stays raw: the callback runs
+    /// user JS, which can re-enter this function's dispatch through a nested
+    /// event loop.
+    ///
+    /// # Safety
+    /// `this` is a live threadsafe function on the JS thread.
+    unsafe fn call(this: *mut Self, task: *mut c_void) {
+        // SAFETY: scoped reborrow; `env` is a raw pointer copy.
+        let Some(env) = (unsafe { &(*this).env }).as_ref().map(NapiEnvRef::get) else {
             // env torn down; nothing to call into.
-            return Ok(());
+            return;
         };
-        if !is_first {
-            let Some(loop_) = self.loop_mut() else {
-                return Ok(());
-            };
-            loop_.drain_microtasks()?;
+
+        // Copy the callback target out so no borrow of `*this` is live while
+        // user code runs.
+        enum Target {
+            Js(JSValue),
+            C(
+                napi_threadsafe_function_call_js,
+                Option<JSValue>,
+                *mut c_void,
+            ),
         }
+        // SAFETY: scoped reborrow.
+        let target = match unsafe { &(*this).callback } {
+            TsfnCallback::Js(strong) => Target::Js(strong.get().unwrap_or(JSValue::UNDEFINED)),
+            TsfnCallback::C {
+                js: cb_js,
+                napi_threadsafe_function_call_js,
+            } => Target::C(
+                *napi_threadsafe_function_call_js,
+                cb_js.get(),
+                // SAFETY: plain field read, same reborrow scope.
+                unsafe { (*this).ctx },
+            ),
+        };
+
         // SAFETY: env is valid while the TSF is live.
         let global_object = unsafe { &*env }.to_js();
 
-        let _dispatch = self.tracker.dispatch(global_object);
+        // SAFETY: scoped reborrow; `tracker` is `Copy`, the guard borrows only
+        // `global_object`.
+        let _dispatch = unsafe { (*this).tracker }.dispatch(global_object);
 
-        match &self.callback {
-            TsfnCallback::Js(strong) => {
-                let js: JSValue = strong.get().unwrap_or(JSValue::UNDEFINED);
+        match target {
+            Target::Js(js) => {
                 if js.is_empty_or_undefined_or_null() {
-                    return Ok(());
+                    return;
                 }
 
                 let _ = js
                     .call(global_object, JSValue::UNDEFINED, &[])
                     .map_err(|err| global_object.report_active_exception_as_unhandled(err));
             }
-            TsfnCallback::C {
-                js: cb_js,
-                napi_threadsafe_function_call_js,
-            } => {
-                // SAFETY: `env` is held alive by `self.env` (`NapiEnvRef`) for the TSF's lifetime.
+            Target::C(call_js, cb_js, ctx) => {
+                // SAFETY: `env` is held alive by `(*this).env` (`NapiEnvRef`) for the TSF's lifetime.
                 let env_ref = unsafe { &*env };
                 let _hs = NapiHandleScope::open_scoped(env_ref);
                 // No func at creation => null js_callback (Node), not encoded undefined.
-                let js = match cb_js.get() {
+                let js = match cb_js {
                     Some(v) => napi_value::create(env_ref, v),
                     None => napi_value(0),
                 };
-                napi_threadsafe_function_call_js(env, js, self.ctx, task);
+                call_js(env, js, ctx, task);
             }
         }
-        Ok(())
     }
 
     /// Runs on an addon thread. A call that reports `napi_closing` consumes the
@@ -2780,22 +2870,24 @@
         let prev = self
             .dispatch_state
             .swap(DispatchState::Pending as u8, Ordering::SeqCst);
-        match prev {
-            x if x == DispatchState::Idle as u8 => {
-                let self_ptr: *mut Self = self;
-                let Some(event_loop) = self.event_loop.as_ref() else {
-                    // env torn down: the loop is gone, nothing to schedule onto.
-                    return;
-                };
-                event_loop.enqueue_task_concurrent(ConcurrentTask::create_from(self_ptr));
-            }
-            x if x == DispatchState::Running as u8 => {
-                // it will check if it has more work to do
-            }
-            _ => {
-                // we've already scheduled it to run
-            }
+        if prev == DispatchState::Pending as u8 {
+            // we've already scheduled it to run
+            return;
         }
+        if prev == DispatchState::Running as u8
+            && self.inflight_dispatch_tasks.load(Ordering::SeqCst) > 0
+        {
+            // the running loop, or the already-queued task if it blocks in a
+            // nested event loop (#36828), picks the item up
+            return;
+        }
+        let self_ptr: *mut Self = self;
+        let Some(event_loop) = self.event_loop.as_ref() else {
+            // env torn down: the loop is gone, nothing to schedule onto.
+            return;
+        };
+        let _ = self.inflight_dispatch_tasks.fetch_add(1, Ordering::SeqCst);
+        event_loop.enqueue_task_concurrent(ConcurrentTask::create_from(self_ptr));
     }
 
     /// Consumes and frees a heap-allocated ThreadSafeFunction (allocated by `new`).
@@ -3082,6 +3174,9 @@
         has_queued_finalizer: false,
         lock: Mutex::new(),
         dispatch_state: AtomicU8::new(DispatchState::Idle as u8),
+        dispatch_depth: 0,
+        pending_destroy: false,
+        inflight_dispatch_tasks: AtomicU32::new(0),
         blocking_condvar: Condvar::default(),
         closing: AtomicU8::new(ClosingState::NotClosing as u8),
         env_dead: AtomicBool::new(false),