grpc 1.57.0 → 1.58.0.pre1

data/src/core/lib/event_engine/thread_pool/thread_count.h

@@ -0,0 +1,161 @@
+// Copyright 2023 The gRPC Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#ifndef GRPC_SRC_CORE_LIB_EVENT_ENGINE_THREAD_POOL_THREAD_COUNT_H
+#define GRPC_SRC_CORE_LIB_EVENT_ENGINE_THREAD_POOL_THREAD_COUNT_H
+
+#include <grpc/support/port_platform.h>
+
+#include <atomic>
+#include <cstddef>
+#include <numeric>
+#include <utility>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+
+#include <grpc/support/cpu.h>
+
+#include "src/core/lib/gpr/useful.h"
+#include "src/core/lib/gprpp/sync.h"
+#include "src/core/lib/gprpp/time.h"
+
+namespace grpc_event_engine {
+namespace experimental {
+
+// Tracks counts across some fixed number of shards.
+// It is intended for fast increment/decrement operations, but a slower overall
+// count operation.
+class BusyThreadCount {
+ public:
+  // Increments a per-shard counter on construction, decrements on destruction.
+  class AutoThreadCounter {
+   public:
+    AutoThreadCounter(BusyThreadCount* counter, size_t idx)
+        : counter_(counter), idx_(idx) {
+      counter_->Increment(idx_);
+    }
+    ~AutoThreadCounter() {
+      if (counter_ != nullptr) counter_->Decrement(idx_);
+    }
+    // not copyable
+    AutoThreadCounter(const AutoThreadCounter&) = delete;
+    AutoThreadCounter& operator=(const AutoThreadCounter&) = delete;
+    // moveable
+    AutoThreadCounter(AutoThreadCounter&& other) noexcept {
+      counter_ = std::exchange(other.counter_, nullptr);
+      idx_ = other.idx_;
+    }
+    AutoThreadCounter& operator=(AutoThreadCounter&& other) noexcept {
+      counter_ = std::exchange(other.counter_, nullptr);
+      idx_ = other.idx_;
+      return *this;
+    }
+
+   private:
+    BusyThreadCount* counter_;
+    size_t idx_;
+  };
+
+  BusyThreadCount() : shards_(grpc_core::Clamp(gpr_cpu_num_cores(), 2u, 64u)) {}
+  AutoThreadCounter MakeAutoThreadCounter(size_t idx) {
+    return AutoThreadCounter(this, idx);
+  };
+  void Increment(size_t idx) {
+    shards_[idx].busy_count.fetch_add(1, std::memory_order_relaxed);
+  }
+  void Decrement(size_t idx) {
+    shards_[idx].busy_count.fetch_sub(1, std::memory_order_relaxed);
+  }
+  size_t count() {
+    return std::accumulate(
+        shards_.begin(), shards_.end(), 0, [](size_t running, ShardedData& d) {
+          return running + d.busy_count.load(std::memory_order_relaxed);
+        });
+  }
+  // Returns some valid index into the per-shard data, which is rotated on every
+  // call to distribute load and reduce contention.
+  size_t NextIndex() { return next_idx_.fetch_add(1) % shards_.size(); }
+
+ private:
+  struct ShardedData {
+    std::atomic<size_t> busy_count{0};
+  } GPR_ALIGN_STRUCT(GPR_CACHELINE_SIZE);
+
+  std::vector<ShardedData> shards_;
+  std::atomic<size_t> next_idx_{0};
+};
+
+// Tracks the number of living threads. It is intended for a fast count
+// operation, with relatively slower increment/decrement operations.
+class LivingThreadCount {
+ public:
+  // Increments the global counter on construction, decrements on destruction.
+  class AutoThreadCounter {
+   public:
+    explicit AutoThreadCounter(LivingThreadCount* counter) : counter_(counter) {
+      counter_->Increment();
+    }
+    ~AutoThreadCounter() {
+      if (counter_ != nullptr) counter_->Decrement();
+    }
+    // not copyable
+    AutoThreadCounter(const AutoThreadCounter&) = delete;
+    AutoThreadCounter& operator=(const AutoThreadCounter&) = delete;
+    // moveable
+    AutoThreadCounter(AutoThreadCounter&& other) noexcept {
+      counter_ = std::exchange(other.counter_, nullptr);
+    }
+    AutoThreadCounter& operator=(AutoThreadCounter&& other) noexcept {
+      counter_ = std::exchange(other.counter_, nullptr);
+      return *this;
+    }
+
+   private:
+    LivingThreadCount* counter_;
+  };
+
+  AutoThreadCounter MakeAutoThreadCounter() { return AutoThreadCounter(this); };
+  void Increment() ABSL_LOCKS_EXCLUDED(mu_) {
+    grpc_core::MutexLock lock(&mu_);
+    ++living_count_;
+    cv_.SignalAll();
+  }
+  void Decrement() ABSL_LOCKS_EXCLUDED(mu_) {
+    grpc_core::MutexLock lock(&mu_);
+    --living_count_;
+    cv_.SignalAll();
+  }
+  void BlockUntilThreadCount(size_t desired_threads, const char* why)
+      ABSL_LOCKS_EXCLUDED(mu_);
+  size_t count() ABSL_LOCKS_EXCLUDED(mu_) {
+    grpc_core::MutexLock lock(&mu_);
+    return CountLocked();
+  }
+
+ private:
+  size_t CountLocked() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    return living_count_;
+  }
+  size_t WaitForCountChange(size_t desired_threads,
+                            grpc_core::Duration timeout);
+
+  grpc_core::Mutex mu_;
+  grpc_core::CondVar cv_ ABSL_GUARDED_BY(mu_);
+  size_t living_count_ ABSL_GUARDED_BY(mu_) = 0;
+};
+
+}  // namespace experimental
+}  // namespace grpc_event_engine
+
+#endif  // GRPC_SRC_CORE_LIB_EVENT_ENGINE_THREAD_POOL_THREAD_COUNT_H

data/src/core/lib/event_engine/thread_pool/thread_pool_factory.cc

@@ -29,6 +29,13 @@ namespace grpc_event_engine {
 namespace experimental {
 
 std::shared_ptr<ThreadPool> MakeThreadPool(size_t reserve_threads) {
+// TODO(hork): remove when the listener flake is identified
+#ifdef GPR_WINDOWS
+  if (grpc_core::IsEventEngineListenerEnabled()) {
+    return std::make_shared<WorkStealingThreadPool>(
+        grpc_core::Clamp(gpr_cpu_num_cores(), 2u, 16u));
+  }
+#endif
   if (grpc_core::IsWorkStealingEnabled()) {
     return std::make_shared<WorkStealingThreadPool>(
         grpc_core::Clamp(gpr_cpu_num_cores(), 2u, 16u));

data/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.cc

@@ -23,7 +23,6 @@
 #include <inttypes.h>
 
 #include <atomic>
-#include <cmath>
 #include <memory>
 #include <utility>
 
@@ -39,10 +38,50 @@
 #include "src/core/lib/event_engine/trace.h"
 #include "src/core/lib/event_engine/work_queue/basic_work_queue.h"
 #include "src/core/lib/event_engine/work_queue/work_queue.h"
-#include "src/core/lib/gpr/time_precise.h"
 #include "src/core/lib/gprpp/thd.h"
 #include "src/core/lib/gprpp/time.h"
 
+// ## Thread Pool Fork-handling
+//
+// Thread-safety needs special attention with regard to fork() calls. The
+// Forkable system employs a pre- and post- fork callback system that does not
+// guarantee any ordering of execution. On fork() events, the thread pool does
+// the following:
+//
+// On pre-fork:
+// * the WorkStealingThreadPool triggers all threads to exit,
+// * all queued work is saved, and
+// * all threads will are down, including the Lifeguard thread.
+//
+// On post-fork:
+//  * all threads are restarted, including the Lifeguard thread, and
+//  * all previously-saved work is enqueued for execution.
+//
+// However, the queue may may get into trouble if one thread is attempting to
+// restart the thread pool while another thread is shutting it down. For that
+// reason, Quiesce and Start must be thread-safe, and Quiesce must wait for the
+// pool to be in a fully started state before it is allowed to continue.
+// Consider this potential ordering of events between Start and Quiesce:
+//
+//     ┌──────────┐
+//     │ Thread 1 │
+//     └────┬─────┘  ┌──────────┐
+//          │        │ Thread 2 │
+//          ▼        └────┬─────┘
+//        Start()         │
+//          │             ▼
+//          │        Quiesce()
+//          │        Wait for worker threads to exit
+//          │        Wait for the lifeguard thread to exit
+//          ▼
+//        Start the Lifeguard thread
+//        Start the worker threads
+//
+// Thread 2 will find no worker threads, and it will then want to wait on a
+// non-existent Lifeguard thread to finish. Trying a simple
+// `lifeguard_thread_.Join()` leads to memory access errors. This implementation
+// uses Notifications to coordinate startup and shutdown states.
+
 namespace grpc_event_engine {
 namespace experimental {
 
@@ -141,9 +180,11 @@ void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Run(
   GPR_DEBUG_ASSERT(quiesced_.load(std::memory_order_relaxed) == false);
   if (g_local_queue != nullptr) {
     g_local_queue->Add(closure);
-    return;
+  } else {
+    queue_.Add(closure);
   }
-  queue_.Add(closure);
+  // Signal a worker in any case, even if work was added to a local queue. This
+  // improves performance on 32-core streaming benchmarks with small payloads.
   work_signal_.Signal();
 }
 
@@ -169,16 +210,13 @@ void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Quiesce() {
   // Note that if this is a threadpool thread then we won't exit this thread
   // until all other threads have exited, so we need to wait for just one thread
   // running instead of zero.
-  gpr_cycle_counter start_time = gpr_get_cycle_counter();
   bool is_threadpool_thread = g_local_queue != nullptr;
-  thread_count()->BlockUntilThreadCount(CounterType::kLivingThreadCount,
-                                        is_threadpool_thread ? 1 : 0,
-                                        "shutting down", work_signal());
+  work_signal()->SignalAll();
+  living_thread_count_.BlockUntilThreadCount(is_threadpool_thread ? 1 : 0,
+                                             "shutting down");
   GPR_ASSERT(queue_.Empty());
   quiesced_.store(true, std::memory_order_relaxed);
-  lifeguard_.BlockUntilShutdown();
-  GRPC_EVENT_ENGINE_TRACE("%ld cycles spent quiescing the pool",
-                          std::lround(gpr_get_cycle_counter() - start_time));
+  lifeguard_.BlockUntilShutdownAndReset();
 }
 
 bool WorkStealingThreadPool::WorkStealingThreadPoolImpl::SetThrottled(
@@ -213,9 +251,9 @@ bool WorkStealingThreadPool::WorkStealingThreadPoolImpl::IsQuiesced() {
 
 void WorkStealingThreadPool::WorkStealingThreadPoolImpl::PrepareFork() {
   SetForking(true);
-  thread_count()->BlockUntilThreadCount(CounterType::kLivingThreadCount, 0,
-                                        "forking", &work_signal_);
-  lifeguard_.BlockUntilShutdown();
+  work_signal_.SignalAll();
+  living_thread_count_.BlockUntilThreadCount(0, "forking");
+  lifeguard_.BlockUntilShutdownAndReset();
 }
 
 void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Postfork() {
@@ -231,13 +269,14 @@ WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::Lifeguard(
       backoff_(grpc_core::BackOff::Options()
                    .set_initial_backoff(kLifeguardMinSleepBetweenChecks)
                    .set_max_backoff(kLifeguardMaxSleepBetweenChecks)
-                   .set_multiplier(1.3)) {}
+                   .set_multiplier(1.3)),
+      lifeguard_should_shut_down_(std::make_unique<grpc_core::Notification>()),
+      lifeguard_is_shut_down_(std::make_unique<grpc_core::Notification>()) {}
 
 void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::Start() {
   // lifeguard_running_ is set early to avoid a quiesce race while the
   // lifeguard is still starting up.
-  grpc_core::MutexLock lock(&lifeguard_shutdown_mu_);
-  lifeguard_running_ = true;
+  lifeguard_running_.store(true);
   grpc_core::Thread(
       "lifeguard",
       [](void* arg) {
@@ -251,7 +290,6 @@ void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::Start() {
 
 void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::
     LifeguardMain() {
-  grpc_core::MutexLock lock(&lifeguard_shutdown_mu_);
   while (true) {
     if (pool_->IsForking()) break;
     // If the pool is shut down, loop quickly until quiesced. Otherwise,
@@ -259,29 +297,33 @@ void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::
     if (pool_->IsShutdown()) {
       if (pool_->IsQuiesced()) break;
     } else {
-      lifeguard_shutdown_cv_.WaitWithTimeout(
-          &lifeguard_shutdown_mu_,
+      lifeguard_should_shut_down_->WaitForNotificationWithTimeout(
           absl::Milliseconds(
               (backoff_.NextAttemptTime() - grpc_core::Timestamp::Now())
                   .millis()));
     }
     MaybeStartNewThread();
   }
-  lifeguard_running_ = false;
-  lifeguard_shutdown_cv_.Signal();
+  lifeguard_running_.store(false, std::memory_order_relaxed);
+  lifeguard_is_shut_down_->Notify();
 }
 
 void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::
-    BlockUntilShutdown() {
-  grpc_core::MutexLock lock(&lifeguard_shutdown_mu_);
-  while (lifeguard_running_) {
-    lifeguard_shutdown_cv_.Signal();
-    lifeguard_shutdown_cv_.WaitWithTimeout(
-        &lifeguard_shutdown_mu_, absl::Seconds(kBlockingQuiesceLogRateSeconds));
+    BlockUntilShutdownAndReset() {
+  lifeguard_should_shut_down_->Notify();
+  while (lifeguard_running_.load(std::memory_order_relaxed)) {
     GRPC_LOG_EVERY_N_SEC_DELAYED(kBlockingQuiesceLogRateSeconds, GPR_DEBUG,
                                  "%s",
                                  "Waiting for lifeguard thread to shut down");
+    lifeguard_is_shut_down_->WaitForNotification();
   }
+  // Do an additional wait in case this method races with LifeguardMain's
+  // shutdown. This should return immediately if the lifeguard is already shut
+  // down.
+  lifeguard_is_shut_down_->WaitForNotification();
+  backoff_.Reset();
+  lifeguard_should_shut_down_ = std::make_unique<grpc_core::Notification>();
+  lifeguard_is_shut_down_ = std::make_unique<grpc_core::Notification>();
 }
 
 void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::
@@ -289,12 +331,9 @@ void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::
   // No new threads are started when forking.
   // No new work is done when forking needs to begin.
   if (pool_->forking_.load()) return;
-  int busy_thread_count =
-      pool_->thread_count_.GetCount(CounterType::kBusyCount);
-  int living_thread_count =
-      pool_->thread_count_.GetCount(CounterType::kLivingThreadCount);
+  const auto living_thread_count = pool_->living_thread_count()->count();
   // Wake an idle worker thread if there's global work to be had.
-  if (busy_thread_count < living_thread_count) {
+  if (pool_->busy_thread_count()->count() < living_thread_count) {
     if (!pool_->queue_.Empty()) {
       pool_->work_signal()->Signal();
       backoff_.Reset();
@@ -317,7 +356,8 @@ void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::
   // queue, nor any work to steal. Add more sophisticated logic about when to
   // start a thread.
   GRPC_EVENT_ENGINE_TRACE(
-      "Starting new ThreadPool thread due to backlog (total threads: %d)",
+      "Starting new ThreadPool thread due to backlog (total threads: %" PRIuPTR
+      ")",
       living_thread_count + 1);
   pool_->StartThread();
   // Tell the lifeguard to monitor the pool more closely.
@@ -329,12 +369,13 @@ void WorkStealingThreadPool::WorkStealingThreadPoolImpl::Lifeguard::
 WorkStealingThreadPool::ThreadState::ThreadState(
     std::shared_ptr<WorkStealingThreadPoolImpl> pool)
     : pool_(std::move(pool)),
-      auto_thread_count_(pool_->thread_count(),
-                         CounterType::kLivingThreadCount),
+      auto_thread_counter_(
+          pool_->living_thread_count()->MakeAutoThreadCounter()),
       backoff_(grpc_core::BackOff::Options()
                    .set_initial_backoff(kWorkerThreadMinSleepBetweenChecks)
                    .set_max_backoff(kWorkerThreadMaxSleepBetweenChecks)
-                   .set_multiplier(1.3)) {}
+                   .set_multiplier(1.3)),
+      busy_count_idx_(pool_->busy_thread_count()->NextIndex()) {}
 
 void WorkStealingThreadPool::ThreadState::ThreadBody() {
   g_local_queue = new BasicWorkQueue();
@@ -372,8 +413,8 @@ bool WorkStealingThreadPool::ThreadState::Step() {
   auto* closure = g_local_queue->PopMostRecent();
   // If local work is available, run it.
   if (closure != nullptr) {
-    ThreadCount::AutoThreadCount auto_busy{pool_->thread_count(),
-                                           CounterType::kBusyCount};
+    auto busy =
+        pool_->busy_thread_count()->MakeAutoThreadCounter(busy_count_idx_);
     closure->Run();
     return true;
   }
@@ -406,11 +447,11 @@ bool WorkStealingThreadPool::ThreadState::Step() {
     if (pool_->IsShutdown()) break;
     bool timed_out = pool_->work_signal()->WaitWithTimeout(
         backoff_.NextAttemptTime() - grpc_core::Timestamp::Now());
+    if (pool_->IsForking() || pool_->IsShutdown()) break;
     // Quit a thread if the pool has more than it requires, and this thread
     // has been idle long enough.
     if (timed_out &&
-        pool_->thread_count()->GetCount(CounterType::kLivingThreadCount) >
-            pool_->reserve_threads() &&
+        pool_->living_thread_count()->count() > pool_->reserve_threads() &&
         grpc_core::Timestamp::Now() - start_time > kIdleThreadLimit) {
       return false;
     }
@@ -421,8 +462,8 @@ bool WorkStealingThreadPool::ThreadState::Step() {
     return false;
   }
   if (closure != nullptr) {
-    ThreadCount::AutoThreadCount auto_busy{pool_->thread_count(),
-                                           CounterType::kBusyCount};
+    auto busy =
+        pool_->busy_thread_count()->MakeAutoThreadCounter(busy_count_idx_);
     closure->Run();
   }
   backoff_.Reset();
@@ -431,8 +472,8 @@ bool WorkStealingThreadPool::ThreadState::Step() {
 
 void WorkStealingThreadPool::ThreadState::FinishDraining() {
   // The thread is definitionally busy while draining
-  ThreadCount::AutoThreadCount auto_busy{pool_->thread_count(),
-                                         CounterType::kBusyCount};
+  auto busy =
+      pool_->busy_thread_count()->MakeAutoThreadCounter(busy_count_idx_);
   // If a fork occurs at any point during shutdown, quit draining. The post-fork
   // threads will finish draining the global queue.
   while (!pool_->IsForking()) {
@@ -454,72 +495,6 @@ void WorkStealingThreadPool::ThreadState::FinishDraining() {
   }
 }
 
-// -------- WorkStealingThreadPool::ThreadCount --------
-
-void WorkStealingThreadPool::ThreadCount::Add(CounterType counter_type) {
-  grpc_core::MutexLock lock(&wait_mu_[counter_type]);
-  ++thread_counts_[counter_type];
-  wait_cv_[counter_type].SignalAll();
-}
-
-void WorkStealingThreadPool::ThreadCount::Remove(CounterType counter_type) {
-  grpc_core::MutexLock lock(&wait_mu_[counter_type]);
-  --thread_counts_[counter_type];
-  wait_cv_[counter_type].SignalAll();
-}
-
-void WorkStealingThreadPool::ThreadCount::BlockUntilThreadCount(
-    CounterType counter_type, size_t desired_threads, const char* why,
-    WorkSignal* work_signal) {
-  // Wait for all threads to exit.
-  while (true) {
-    auto curr_threads = WaitForCountChange(
-        counter_type, desired_threads,
-        grpc_core::Duration::Seconds(kBlockingQuiesceLogRateSeconds));
-    if (curr_threads == desired_threads) break;
-    GRPC_LOG_EVERY_N_SEC_DELAYED(
-        kBlockingQuiesceLogRateSeconds, GPR_DEBUG,
-        "Waiting for thread pool to idle before %s. (%" PRIdPTR " to %" PRIdPTR
-        ")",
-        why, curr_threads, desired_threads);
-    work_signal->SignalAll();
-  }
-}
-
-size_t WorkStealingThreadPool::ThreadCount::WaitForCountChange(
-    CounterType counter_type, size_t desired_threads,
-    grpc_core::Duration timeout) {
-  size_t count;
-  auto deadline = absl::Now() + absl::Milliseconds(timeout.millis());
-  do {
-    grpc_core::MutexLock lock(&wait_mu_[counter_type]);
-    count = GetCountLocked(counter_type);
-    if (count == desired_threads) break;
-    wait_cv_[counter_type].WaitWithDeadline(&wait_mu_[counter_type], deadline);
-  } while (absl::Now() < deadline);
-  return count;
-}
-
-size_t WorkStealingThreadPool::ThreadCount::GetCount(CounterType counter_type) {
-  grpc_core::MutexLock lock(&wait_mu_[counter_type]);
-  return GetCountLocked(counter_type);
-}
-
-size_t WorkStealingThreadPool::ThreadCount::GetCountLocked(
-    CounterType counter_type) {
-  return thread_counts_[counter_type];
-}
-
-WorkStealingThreadPool::ThreadCount::AutoThreadCount::AutoThreadCount(
-    ThreadCount* counter, CounterType counter_type)
-    : counter_(counter), counter_type_(counter_type) {
-  counter_->Add(counter_type_);
-}
-
-WorkStealingThreadPool::ThreadCount::AutoThreadCount::~AutoThreadCount() {
-  counter_->Remove(counter_type_);
-}
-
 // -------- WorkStealingThreadPool::WorkSignal --------
 
 void WorkStealingThreadPool::WorkSignal::Signal() {

data/src/core/lib/event_engine/thread_pool/work_stealing_thread_pool.h

@@ -33,9 +33,11 @@
 #include <grpc/event_engine/event_engine.h>
 
 #include "src/core/lib/backoff/backoff.h"
+#include "src/core/lib/event_engine/thread_pool/thread_count.h"
 #include "src/core/lib/event_engine/thread_pool/thread_pool.h"
 #include "src/core/lib/event_engine/work_queue/basic_work_queue.h"
 #include "src/core/lib/event_engine/work_queue/work_queue.h"
+#include "src/core/lib/gprpp/notification.h"
 #include "src/core/lib/gprpp/sync.h"
 #include "src/core/lib/gprpp/time.h"
 
@@ -75,56 +77,6 @@ class WorkStealingThreadPool final : public ThreadPool {
     grpc_core::CondVar cv_ ABSL_GUARDED_BY(mu_);
   };
 
-  // Types of thread counts.
-  // Note this is intentionally not an enum class, the keys are used as indexes
-  // into the ThreadCount's private array.
-  enum CounterType {
-    kLivingThreadCount = 0,
-    kBusyCount,
-  };
-
-  class ThreadCount {
-   public:
-    // Adds 1 to the thread count for that counter type.
-    void Add(CounterType counter_type)
-        ABSL_LOCKS_EXCLUDED(wait_mu_[counter_type]);
-    // Subtracts 1 from the thread count for that counter type.
-    void Remove(CounterType counter_type)
-        ABSL_LOCKS_EXCLUDED(wait_mu_[counter_type]);
-    // Blocks until the thread count for that type reaches `desired_threads`.
-    void BlockUntilThreadCount(CounterType counter_type, size_t desired_threads,
-                               const char* why, WorkSignal* work_signal)
-        ABSL_LOCKS_EXCLUDED(wait_mu_[counter_type]);
-    // Returns the current thread count for the tracked type.
-    size_t GetCount(CounterType counter_type)
-        ABSL_LOCKS_EXCLUDED(wait_mu_[counter_type]);
-    // Returns the current thread count for the tracked type.
-    size_t GetCountLocked(CounterType counter_type)
-        ABSL_EXCLUSIVE_LOCKS_REQUIRED(wait_mu_[counter_type]);
-
-    // Adds and removes thread counts on construction and destruction
-    class AutoThreadCount {
-     public:
-      AutoThreadCount(ThreadCount* counter, CounterType counter_type);
-      ~AutoThreadCount();
-
-     private:
-      ThreadCount* counter_;
-      CounterType counter_type_;
-    };
-
-   private:
-    // Wait for the desired count to be reached.
-    // Returns the current thread count either when the desired count is
-    // reached, or when the deadline has passed, whichever happens first.
-    size_t WaitForCountChange(CounterType counter_type, size_t desired_threads,
-                              grpc_core::Duration timeout);
-
-    grpc_core::Mutex wait_mu_[2];
-    grpc_core::CondVar wait_cv_[2];
-    size_t thread_counts_[2]{0, 0};
-  };
-
   // A pool of WorkQueues that participate in work stealing.
   //
   // Every worker thread registers and unregisters its thread-local thread pool
@@ -185,7 +137,8 @@ class WorkStealingThreadPool final : public ThreadPool {
     bool IsForking();
     bool IsQuiesced();
     size_t reserve_threads() { return reserve_threads_; }
-    ThreadCount* thread_count() { return &thread_count_; }
+    BusyThreadCount* busy_thread_count() { return &busy_thread_count_; }
+    LivingThreadCount* living_thread_count() { return &living_thread_count_; }
     TheftRegistry* theft_registry() { return &theft_registry_; }
     WorkQueue* queue() { return &queue_; }
     WorkSignal* work_signal() { return &work_signal_; }
@@ -202,7 +155,8 @@ class WorkStealingThreadPool final : public ThreadPool {
       // Start the lifeguard thread.
       void Start();
       // Block until the lifeguard thread is shut down.
-      void BlockUntilShutdown();
+      // Afterwards, reset the lifeguard state so it can start again cleanly.
+      void BlockUntilShutdownAndReset();
 
      private:
       // The main body of the lifeguard thread.
@@ -213,14 +167,14 @@ class WorkStealingThreadPool final : public ThreadPool {
       WorkStealingThreadPoolImpl* pool_;
       grpc_core::BackOff backoff_;
       // Used for signaling that the lifeguard thread has stopped running.
-      grpc_core::Mutex lifeguard_shutdown_mu_;
-      bool lifeguard_running_ ABSL_GUARDED_BY(lifeguard_shutdown_mu_) = false;
-      grpc_core::CondVar lifeguard_shutdown_cv_
-          ABSL_GUARDED_BY(lifeguard_shutdown_mu_);
+      std::unique_ptr<grpc_core::Notification> lifeguard_should_shut_down_;
+      std::unique_ptr<grpc_core::Notification> lifeguard_is_shut_down_;
+      std::atomic<bool> lifeguard_running_{false};
     };
 
     const size_t reserve_threads_;
-    ThreadCount thread_count_;
+    BusyThreadCount busy_thread_count_;
+    LivingThreadCount living_thread_count_;
     TheftRegistry theft_registry_;
     BasicWorkQueue queue_;
     // Track shutdown and fork bits separately.
@@ -253,11 +207,11 @@ class WorkStealingThreadPool final : public ThreadPool {
     // is decremented at time of destruction. This is necessary when this thread
     // state holds the last shared_ptr keeping the pool alive.
     std::shared_ptr<WorkStealingThreadPoolImpl> pool_;
-    // auto_thread_count_ must be the second member declared, so that the thread
-    // count is decremented after all other state is cleaned up (preventing
-    // leaks).
-    ThreadCount::AutoThreadCount auto_thread_count_;
+    // auto_thread_counter_ must be declared after pool_, so that the thread
+    // count is decremented after all other pool state is cleaned up.
+    LivingThreadCount::AutoThreadCounter auto_thread_counter_;
     grpc_core::BackOff backoff_;
+    size_t busy_count_idx_;
   };
 
   const std::shared_ptr<WorkStealingThreadPoolImpl> pool_;

data/src/core/lib/event_engine/windows/windows_endpoint.cc

@@ -284,6 +284,7 @@ void WindowsEndpoint::HandleWriteClosure::Prime(
 
 void WindowsEndpoint::HandleReadClosure::Run() {
   // Deletes the shared_ptr when this closure returns
+  // Note that the endpoint may have already been destroyed.
   auto io_state = std::move(io_state_);
   GRPC_EVENT_ENGINE_ENDPOINT_TRACE("WindowsEndpoint::%p Handling Read Event",
                                    io_state->endpoint);
@@ -297,10 +298,8 @@ void WindowsEndpoint::HandleReadClosure::Run() {
   if (result.bytes_transferred == 0) {
     // Either the endpoint is shut down or we've seen the end of the stream
     if (grpc_event_engine_endpoint_data_trace.enabled()) {
-      DumpSliceBuffer(
-          buffer_, absl::StrFormat("WindowsEndpoint::%p READ (peer=%s)",
-                                   io_state->endpoint,
-                                   io_state->endpoint->peer_address_string_));
+      DumpSliceBuffer(buffer_, absl::StrFormat("WindowsEndpoint::%p READ",
+                                               io_state->endpoint));
     }
     status = absl::UnavailableError("End of TCP stream");
     grpc_core::StatusSetInt(&status, grpc_core::StatusIntProperty::kRpcStatus,