grpc 1.54.3 → 1.59.2

data/src/core/lib/event_engine/posix_engine/posix_engine_listener_utils.cc

@@ -148,7 +148,8 @@ absl::Status PrepareSocket(const PosixTcpOptions& options,
     }
   });
   if (PosixSocketWrapper::IsSocketReusePortSupported() &&
-      options.allow_reuse_port && socket.addr.address()->sa_family != AF_UNIX) {
+      options.allow_reuse_port && socket.addr.address()->sa_family != AF_UNIX &&
+      !ResolvedAddressIsVSock(socket.addr)) {
     GRPC_RETURN_IF_ERROR(socket.sock.SetSocketReusePort(1));
   }
 
@@ -164,7 +165,8 @@ absl::Status PrepareSocket(const PosixTcpOptions& options,
   GRPC_RETURN_IF_ERROR(socket.sock.SetSocketNonBlocking(1));
   GRPC_RETURN_IF_ERROR(socket.sock.SetSocketCloexec(1));
 
-  if (socket.addr.address()->sa_family != AF_UNIX) {
+  if (socket.addr.address()->sa_family != AF_UNIX &&
+      !ResolvedAddressIsVSock(socket.addr)) {
     GRPC_RETURN_IF_ERROR(socket.sock.SetSocketLowLatency(1));
     GRPC_RETURN_IF_ERROR(socket.sock.SetSocketReuseAddr(1));
     socket.sock.TrySetSocketTcpUserTimeout(options, false);

data/src/core/lib/event_engine/posix_engine/tcp_socket_utils.cc

@@ -27,6 +27,7 @@
 #include "absl/types/optional.h"
 
 #include <grpc/event_engine/event_engine.h>
+#include <grpc/impl/channel_arg_names.h>
 
 #include "src/core/lib/gpr/useful.h"
 #include "src/core/lib/gprpp/crash.h"  // IWYU pragma: keep
@@ -125,11 +126,14 @@ absl::Status PrepareTcpClientSocket(PosixSocketWrapper sock,
   });
   GRPC_RETURN_IF_ERROR(sock.SetSocketNonBlocking(1));
   GRPC_RETURN_IF_ERROR(sock.SetSocketCloexec(1));
-
-  if (reinterpret_cast<const sockaddr*>(addr.address())->sa_family != AF_UNIX) {
-    // If its not a unix socket address.
+  if (options.tcp_receive_buffer_size != options.kReadBufferSizeUnset) {
+    GRPC_RETURN_IF_ERROR(sock.SetSocketRcvBuf(options.tcp_receive_buffer_size));
+  }
+  if (addr.address()->sa_family != AF_UNIX && !ResolvedAddressIsVSock(addr)) {
+    // If its not a unix socket or vsock address.
     GRPC_RETURN_IF_ERROR(sock.SetSocketLowLatency(1));
     GRPC_RETURN_IF_ERROR(sock.SetSocketReuseAddr(1));
+    GRPC_RETURN_IF_ERROR(sock.SetSocketDscp(options.dscp));
     sock.TrySetSocketTcpUserTimeout(options, true);
   }
   GRPC_RETURN_IF_ERROR(sock.SetSocketNoSigpipeIfPossible());
@@ -169,6 +173,9 @@ PosixTcpOptions TcpOptionsFromEndpointConfig(const EndpointConfig& config) {
   options.tcp_tx_zerocopy_max_simultaneous_sends =
       AdjustValue(PosixTcpOptions::kDefaultMaxSends, 0, INT_MAX,
                   config.GetInt(GRPC_ARG_TCP_TX_ZEROCOPY_MAX_SIMULT_SENDS));
+  options.tcp_receive_buffer_size =
+      AdjustValue(PosixTcpOptions::kReadBufferSizeUnset, 0, INT_MAX,
+                  config.GetInt(GRPC_ARG_TCP_RECEIVE_BUFFER_SIZE));
   options.tcp_tx_zero_copy_enabled =
       (AdjustValue(PosixTcpOptions::kZerocpTxEnabledDefault, 0, 1,
                    config.GetInt(GRPC_ARG_TCP_TX_ZEROCOPY_ENABLED)) != 0);
@@ -179,6 +186,8 @@ PosixTcpOptions TcpOptionsFromEndpointConfig(const EndpointConfig& config) {
   options.expand_wildcard_addrs =
       (AdjustValue(0, 1, INT_MAX,
                    config.GetInt(GRPC_ARG_EXPAND_WILDCARD_ADDRS)) != 0);
+  options.dscp = AdjustValue(PosixTcpOptions::kDscpNotSet, 0, 63,
+                             config.GetInt(GRPC_ARG_DSCP));
   options.allow_reuse_port = PosixSocketWrapper::IsSocketReusePortSupported();
   auto allow_reuse_port_value = config.GetInt(GRPC_ARG_ALLOW_REUSEPORT);
   if (allow_reuse_port_value.has_value()) {
@@ -514,6 +523,39 @@ absl::Status PosixSocketWrapper::SetSocketLowLatency(int low_latency) {
   return absl::OkStatus();
 }
 
+// Set Differentiated Services Code Point (DSCP)
+absl::Status PosixSocketWrapper::SetSocketDscp(int dscp) {
+  if (dscp == PosixTcpOptions::kDscpNotSet) {
+    return absl::OkStatus();
+  }
+  // The TOS/TrafficClass byte consists of following bits:
+  // | 7 6 5 4 3 2 | 1 0 |
+  // |    DSCP     | ECN |
+  int newval = dscp << 2;
+  int val;
+  socklen_t intlen = sizeof(val);
+  // Get ECN bits from current IP_TOS value unless IPv6 only
+  if (0 == getsockopt(fd_, IPPROTO_IP, IP_TOS, &val, &intlen)) {
+    newval |= (val & 0x3);
+    if (0 != setsockopt(fd_, IPPROTO_IP, IP_TOS, &newval, sizeof(newval))) {
+      return absl::Status(
+          absl::StatusCode::kInternal,
+          absl::StrCat("setsockopt(IP_TOS): ", grpc_core::StrError(errno)));
+    }
+  }
+  // Get ECN from current Traffic Class value if IPv6 is available
+  if (0 == getsockopt(fd_, IPPROTO_IPV6, IPV6_TCLASS, &val, &intlen)) {
+    newval |= (val & 0x3);
+    if (0 !=
+        setsockopt(fd_, IPPROTO_IPV6, IPV6_TCLASS, &newval, sizeof(newval))) {
+      return absl::Status(absl::StatusCode::kInternal,
+                          absl::StrCat("setsockopt(IPV6_TCLASS): ",
+                                       grpc_core::StrError(errno)));
+    }
+  }
+  return absl::OkStatus();
+}
+
 #if GPR_LINUX == 1
 // For Linux, it will be detected to support TCP_USER_TIMEOUT
 #ifndef TCP_USER_TIMEOUT
@@ -707,7 +749,7 @@ absl::StatusOr<PosixSocketWrapper> PosixSocketWrapper::CreateDualStackSocket(
     }
     // If this isn't an IPv4 address, then return whatever we've got.
     if (!ResolvedAddressIsV4Mapped(addr, nullptr)) {
-      if (newfd <= 0) {
+      if (newfd < 0) {
         return ErrorForFd(newfd, addr);
       }
       dsmode = PosixSocketWrapper::DSMode::DSMODE_IPV6;
@@ -794,6 +836,10 @@ absl::Status PosixSocketWrapper::SetSocketReusePort(int /*reuse*/) {
   grpc_core::Crash("unimplemented");
 }
 
+absl::Status PosixSocketWrapper::SetSocketDscp(int /*dscp*/) {
+  grpc_core::Crash("unimplemented");
+}
+
 void PosixSocketWrapper::ConfigureDefaultTcpUserTimeout(bool /*enable*/,
                                                         int /*timeout*/,
                                                         bool /*is_client*/) {}

data/src/core/lib/event_engine/posix_engine/tcp_socket_utils.h

@@ -58,16 +58,21 @@ struct PosixTcpOptions {
   static constexpr int kMaxChunkSize = 32 * 1024 * 1024;
   static constexpr int kDefaultMaxSends = 4;
   static constexpr size_t kDefaultSendBytesThreshold = 16 * 1024;
+  // Let the system decide the proper buffer size.
+  static constexpr int kReadBufferSizeUnset = -1;
+  static constexpr int kDscpNotSet = -1;
   int tcp_read_chunk_size = kDefaultReadChunkSize;
   int tcp_min_read_chunk_size = kDefaultMinReadChunksize;
   int tcp_max_read_chunk_size = kDefaultMaxReadChunksize;
   int tcp_tx_zerocopy_send_bytes_threshold = kDefaultSendBytesThreshold;
   int tcp_tx_zerocopy_max_simultaneous_sends = kDefaultMaxSends;
+  int tcp_receive_buffer_size = kReadBufferSizeUnset;
   bool tcp_tx_zero_copy_enabled = kZerocpTxEnabledDefault;
   int keep_alive_time_ms = 0;
   int keep_alive_timeout_ms = 0;
   bool expand_wildcard_addrs = false;
   bool allow_reuse_port = false;
+  int dscp = kDscpNotSet;
   grpc_core::RefCountedPtr<grpc_core::ResourceQuota> resource_quota;
   struct grpc_socket_mutator* socket_mutator = nullptr;
   PosixTcpOptions() = default;
@@ -132,6 +137,7 @@ struct PosixTcpOptions {
     keep_alive_timeout_ms = other.keep_alive_timeout_ms;
     expand_wildcard_addrs = other.expand_wildcard_addrs;
     allow_reuse_port = other.allow_reuse_port;
+    dscp = other.dscp;
   }
 };
 
@@ -179,6 +185,9 @@ class PosixSocketWrapper {
   // Set SO_REUSEPORT
   absl::Status SetSocketReusePort(int reuse);
 
+  // Set Differentiated Services Code Point (DSCP)
+  absl::Status SetSocketDscp(int dscp);
+
   // Override default Tcp user timeout values if necessary.
   void TrySetSocketTcpUserTimeout(const PosixTcpOptions& options,
                                   bool is_client);

data/src/core/lib/event_engine/posix_engine/timer.h

@@ -77,46 +77,19 @@ class TimerList {
   TimerList(const TimerList&) = delete;
   TimerList& operator=(const TimerList&) = delete;
 
-  // Initialize *timer. When expired or canceled, closure will be called with
-  // error set to indicate if it expired (absl::OkStatus()) or was canceled
-  //(absl::CancelledError()). *closure is guaranteed to be called exactly once,
-  // and application code should check the error to determine how it was
-  // invoked. The application callback is also responsible for maintaining
-  // information about when to free up any user-level state. Behavior is
-  // undefined for a deadline of grpc_core::Timestamp::InfFuture().
+  // Initialize a Timer.
+  // When expired, the closure will be run. If the timer is canceled, the
+  // closure will not be run. Behavior is undefined for a deadline of
+  // grpc_core::Timestamp::InfFuture().
   void TimerInit(Timer* timer, grpc_core::Timestamp deadline,
                  experimental::EventEngine::Closure* closure);
 
-  // Note that there is no timer destroy function. This is because the
-  // timer is a one-time occurrence with a guarantee that the callback will
-  // be called exactly once, either at expiration or cancellation. Thus, all
-  // the internal timer event management state is destroyed just before
-  // that callback is invoked. If the user has additional state associated with
-  // the timer, the user is responsible for determining when it is safe to
-  // destroy that state.
-
-  // Cancel an *timer.
-  // There are three cases:
-  // 1. We normally cancel the timer
-  // 2. The timer has already run
-  // 3. We can't cancel the timer because it is "in flight".
-
-  // In all of these cases, the cancellation is still considered successful.
-  // They are essentially distinguished in that the timer_cb will be run
-  // exactly once from either the cancellation (with error
-  // absl::CancelledError()) or from the activation (with error
-  // absl::OkStatus()).
-
-  // Note carefully that the callback function MAY occur in the same callstack
-  // as grpc_timer_cancel. It's expected that most timers will be cancelled
-  // (their primary use is to implement deadlines), and so this code is
-  // optimized such that cancellation costs as little as possible. Making
-  // callbacks run inline matches this aim.
-
-  // Requires: cancel() must happen after init() on a given timer
-  bool TimerCancel(Timer* timer) GRPC_MUST_USE_RESULT;
-
-  // iomgr internal api for dealing with timers
+  // Cancel a Timer.
+  // Returns false if the timer cannot be canceled. This will happen if the
+  // timer has already fired, or if its closure is currently running. The
+  // closure is guaranteed to run eventually if this method returns false.
+  // Otherwise, this returns true, and the closure will not be run.
+  GRPC_MUST_USE_RESULT bool TimerCancel(Timer* timer);
 
   // Check for timers to be run, and return them.
   // Return nullopt if timers could not be checked due to contention with

data/src/core/lib/event_engine/posix_engine/timer_manager.h

@@ -33,7 +33,7 @@
 
 #include "src/core/lib/event_engine/forkable.h"
 #include "src/core/lib/event_engine/posix_engine/timer.h"
-#include "src/core/lib/event_engine/thread_pool.h"
+#include "src/core/lib/event_engine/thread_pool/thread_pool.h"
 #include "src/core/lib/gprpp/notification.h"
 #include "src/core/lib/gprpp/sync.h"
 #include "src/core/lib/gprpp/thd.h"

data/src/core/lib/event_engine/posix_engine/traced_buffer_list.cc

@@ -279,6 +279,8 @@ void TracedBufferList::ProcessTimestamp(struct sock_extended_err* serr,
       elem = elem->next_;
       continue;
     }
+    g_timestamps_callback(elem->arg_, &(elem->ts_),
+                          absl::DeadlineExceededError("Ack timed out"));
     if (prev != nullptr) {
       prev->next_ = elem->next_;
       delete elem;

data/src/core/lib/event_engine/shim.cc

@@ -22,6 +22,12 @@
 namespace grpc_event_engine {
 namespace experimental {
 
+#if defined(GRPC_POSIX_SOCKET_TCP) && !defined(GRPC_CFSTREAM)
+bool g_event_engine_supports_fd = true;
+#else
+bool g_event_engine_supports_fd = false;
+#endif
+
 bool UseEventEngineClient() {
 // TODO(hork, eryu): Adjust the ifdefs accordingly when event engines become
 // available for other platforms.
@@ -29,6 +35,8 @@ bool UseEventEngineClient() {
   return grpc_core::IsEventEngineClientEnabled();
 #elif defined(GPR_WINDOWS)
   return grpc_core::IsEventEngineClientEnabled();
+#elif defined(GRPC_IOS_EVENT_ENGINE_CLIENT)
+  return true;
 #else
   return false;
 #endif
@@ -46,7 +54,7 @@ bool UseEventEngineListener() {
 
 bool EventEngineSupportsFd() {
 #if defined(GRPC_POSIX_SOCKET_TCP) && !defined(GRPC_CFSTREAM)
-  return true;
+  return g_event_engine_supports_fd;
 #else
   return false;
 #endif

data/src/core/lib/event_engine/tcp_socket_utils.cc

@@ -36,6 +36,10 @@
 #include <sys/un.h>
 #endif
 
+#ifdef GRPC_HAVE_VSOCK
+#include <linux/vm_sockets.h>
+#endif
+
 #include <errno.h>
 #include <inttypes.h>
 #include <stdlib.h>
@@ -70,6 +74,10 @@ absl::StatusOr<std::string> GetScheme(
       return "ipv6";
     case AF_UNIX:
       return "unix";
+#ifdef GRPC_HAVE_VSOCK
+    case AF_VSOCK:
+      return "vsock";
+#endif
     default:
       return absl::InvalidArgumentError(
           absl::StrFormat("Unknown sockaddr family: %d",
@@ -92,12 +100,11 @@ absl::StatusOr<std::string> ResolvedAddrToUnixPathIfPossible(
 #else
   int len = resolved_addr->size() - sizeof(unix_addr->sun_family) - 1;
 #endif
-  bool abstract = (len < 0 || unix_addr->sun_path[0] == '\0');
+  if (len <= 0) return "";
   std::string path;
-  if (abstract) {
-    if (len >= 0) {
-      path = std::string(unix_addr->sun_path + 1, len);
-    }
+  if (unix_addr->sun_path[0] == '\0') {
+    // unix-abstract socket processing.
+    path = std::string(unix_addr->sun_path + 1, len);
     path = absl::StrCat(std::string(1, '\0'), path);
   } else {
     size_t maxlen = sizeof(unix_addr->sun_path);
@@ -115,9 +122,9 @@ absl::StatusOr<std::string> ResolvedAddrToUriUnixIfPossible(
   GRPC_RETURN_IF_ERROR(path.status());
   std::string scheme;
   std::string path_string;
-  if (path->at(0) == '\0') {
+  if (!path->empty() && path->at(0) == '\0' && path->length() > 1) {
     scheme = "unix-abstract";
-    path_string = path->length() > 1 ? path->substr(1, std::string::npos) : "";
+    path_string = path->substr(1, std::string::npos);
   } else {
     scheme = "unix";
     path_string = std::move(*path);
@@ -142,6 +149,39 @@ absl::StatusOr<std::string> ResolvedAddrToUriUnixIfPossible(
 }
 #endif
 
+#ifdef GRPC_HAVE_VSOCK
+absl::StatusOr<std::string> ResolvedAddrToVsockPathIfPossible(
+    const EventEngine::ResolvedAddress* resolved_addr) {
+  const sockaddr* addr = resolved_addr->address();
+  if (addr->sa_family != AF_VSOCK) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Socket family is not AF_VSOCK: ", addr->sa_family));
+  }
+  const sockaddr_vm* vm_addr = reinterpret_cast<const sockaddr_vm*>(addr);
+  return absl::StrCat(vm_addr->svm_cid, ":", vm_addr->svm_port);
+}
+
+absl::StatusOr<std::string> ResolvedAddrToUriVsockIfPossible(
+    const EventEngine::ResolvedAddress* resolved_addr) {
+  auto path = ResolvedAddrToVsockPathIfPossible(resolved_addr);
+  absl::StatusOr<grpc_core::URI> uri =
+      grpc_core::URI::Create("vsock", /*authority=*/"", std::move(*path),
+                             /*query_parameter_pairs=*/{}, /*fragment=*/"");
+  if (!uri.ok()) return uri.status();
+  return uri->ToString();
+}
+#else
+absl::StatusOr<std::string> ResolvedAddrToVsockPathIfPossible(
+    const EventEngine::ResolvedAddress* /*resolved_addr*/) {
+  return absl::InvalidArgumentError("VSOCK is not supported.");
+}
+
+absl::StatusOr<std::string> ResolvedAddrToUriVsockIfPossible(
+    const EventEngine::ResolvedAddress* /*resolved_addr*/) {
+  return absl::InvalidArgumentError("VSOCK is not supported.");
+}
+#endif
+
 }  // namespace
 
 bool ResolvedAddressIsV4Mapped(
@@ -232,6 +272,10 @@ int ResolvedAddressGetPort(const EventEngine::ResolvedAddress& resolved_addr) {
 #ifdef GRPC_HAVE_UNIX_SOCKET
     case AF_UNIX:
       return 1;
+#endif
+#ifdef GRPC_HAVE_VSOCK
+    case AF_VSOCK:
+      return 1;
 #endif
     default:
       gpr_log(GPR_ERROR, "Unknown socket family %d in ResolvedAddressGetPort",
@@ -292,6 +336,15 @@ absl::optional<int> ResolvedAddressIsWildcard(
   }
 }
 
+bool ResolvedAddressIsVSock(const EventEngine::ResolvedAddress& resolved_addr) {
+#ifdef GRPC_HAVE_VSOCK
+  return resolved_addr.address()->sa_family == AF_VSOCK;
+#else
+  (void)resolved_addr;
+  return false;
+#endif
+}
+
 absl::StatusOr<std::string> ResolvedAddressToNormalizedString(
     const EventEngine::ResolvedAddress& resolved_addr) {
   EventEngine::ResolvedAddress addr_normalized;
@@ -312,6 +365,10 @@ absl::StatusOr<std::string> ResolvedAddressToString(
   }
 #endif  // GRPC_HAVE_UNIX_SOCKET
 
+  if (ResolvedAddressIsVSock(resolved_addr)) {
+    return ResolvedAddrToVsockPathIfPossible(&resolved_addr);
+  }
+
   const void* ip = nullptr;
   int port = 0;
   uint32_t sin6_scope_id = 0;
@@ -362,6 +419,9 @@ absl::StatusOr<std::string> ResolvedAddressToURI(
   if (*scheme == "unix") {
     return ResolvedAddrToUriUnixIfPossible(&addr);
   }
+  if (*scheme == "vsock") {
+    return ResolvedAddrToUriVsockIfPossible(&addr);
+  }
   auto path = ResolvedAddressToString(addr);
   GRPC_RETURN_IF_ERROR(path.status());
   absl::StatusOr<grpc_core::URI> uri =

data/src/core/lib/event_engine/tcp_socket_utils.h

@@ -61,6 +61,9 @@ void ResolvedAddressSetPort(EventEngine::ResolvedAddress& resolved_addr,
 absl::optional<int> ResolvedAddressIsWildcard(
     const EventEngine::ResolvedAddress& addr);
 
+// Returns true if resolved_addr is an VSOCK address. Otherwise returns false.
+bool ResolvedAddressIsVSock(const EventEngine::ResolvedAddress& resolved_addr);
+
 // Converts a EventEngine::ResolvedAddress into a newly-allocated
 // human-readable string.
 // Currently, only the AF_INET, AF_INET6, and AF_UNIX families are

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

@@ -0,0 +1,58 @@
+// 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.
+#include <grpc/support/port_platform.h>
+
+#include "src/core/lib/event_engine/thread_pool/thread_count.h"
+
+#include <inttypes.h>
+
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+
+#include <grpc/support/log.h>
+
+namespace grpc_event_engine {
+namespace experimental {
+
+// -------- LivingThreadCount --------
+
+void LivingThreadCount::BlockUntilThreadCount(size_t desired_threads,
+                                              const char* why) {
+  constexpr grpc_core::Duration log_rate = grpc_core::Duration::Seconds(3);
+  while (true) {
+    auto curr_threads = WaitForCountChange(desired_threads, log_rate);
+    if (curr_threads == desired_threads) break;
+    GRPC_LOG_EVERY_N_SEC_DELAYED(
+        log_rate.seconds(), GPR_DEBUG,
+        "Waiting for thread pool to idle before %s. (%" PRIdPTR " to %" PRIdPTR
+        ")",
+        why, curr_threads, desired_threads);
+  }
+}
+
+size_t LivingThreadCount::WaitForCountChange(size_t desired_threads,
+                                             grpc_core::Duration timeout) {
+  size_t count;
+  auto deadline = absl::Now() + absl::Milliseconds(timeout.millis());
+  do {
+    grpc_core::MutexLock lock(&mu_);
+    count = CountLocked();
+    if (count == desired_threads) break;
+    cv_.WaitWithDeadline(&mu_, deadline);
+  } while (absl::Now() < deadline);
+  return count;
+}
+
+}  // namespace experimental
+}  // namespace grpc_event_engine

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

@@ -0,0 +1,176 @@
+// 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 <cstdint>
+#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:
+// We want to ensure that this data structure lands on different cachelines per
+// cpu. With C++17 we can do so explicitly with an `alignas` specifier. Prior
+// versions we can at best approximate it by padding the structure. It'll
+// probably work out ok, but it's not guaranteed across allocators.
+// TODO(ctiller): When we move to C++17 delete the duplicate definition.
+#if __cplusplus >= 201703L
+  struct ShardedData {
+    std::atomic<size_t> busy_count{0};
+  } GPR_ALIGN_STRUCT(GPR_CACHELINE_SIZE);
+#else
+  struct ShardedDataHeader {
+    std::atomic<size_t> busy_count{0};
+  };
+  struct ShardedData : public ShardedDataHeader {
+    uint8_t padding[GPR_CACHELINE_SIZE - sizeof(ShardedDataHeader)];
+  };
+#endif
+
+  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