diff --git a/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.cc b/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.cc index 55e8ec7bcf6..fd60eec3ce7 100644 --- a/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.cc +++ b/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.cc @@ -321,6 +321,8 @@ void WorkStealingThreadPool::ThreadState::ThreadBody() { pool_->queue()->Add(closure); } } + } else if (pool_->IsShutdown()) { + FinishDraining(); } GPR_ASSERT(g_local_queue->Empty()); pool_->theft_registry()->Unenroll(g_local_queue); @@ -395,6 +397,28 @@ bool WorkStealingThreadPool::ThreadState::Step() { return should_run_again; } +void WorkStealingThreadPool::ThreadState::FinishDraining() { + // If a fork occurs at any point during shutdown, quit draining. The post-fork + // threads will finish draining the global queue. + while (!pool_->IsForking()) { + if (!g_local_queue->Empty()) { + auto* closure = g_local_queue->PopMostRecent(); + if (closure != nullptr) { + closure->Run(); + } + continue; + } + if (!pool_->queue()->Empty()) { + auto* closure = pool_->queue()->PopMostRecent(); + if (closure != nullptr) { + closure->Run(); + } + continue; + } + break; + } +} + // -------- WorkStealingThreadPool::ThreadCount -------- void WorkStealingThreadPool::ThreadCount::Add(CounterType counter_type) { diff --git a/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.h b/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.h index 14c366e4d1a..f5cbba8debd 100644 --- a/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.h +++ b/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.h @@ -224,6 +224,9 @@ class WorkStealingThreadPool final : public ThreadPool { void ThreadBody(); void SleepIfRunning(); bool Step(); + // After the pool is shut down, ensure all local and global callbacks are + // executed before quitting the thread. + void FinishDraining(); private: // pool_ must be the first member so that it is alive when the thread count