grpc 1.57.0.pre1 → 1.58.0.pre1

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

@@ -0,0 +1,73 @@
+// 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_GRPC_POLLED_FD_H
+#define GRPC_SRC_CORE_LIB_EVENT_ENGINE_GRPC_POLLED_FD_H
+
+#include <grpc/support/port_platform.h>
+
+#if GRPC_ARES == 1
+
+#include <ares.h>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/status/status.h"
+
+#include "src/core/lib/iomgr/error.h"
+
+namespace grpc_event_engine {
+namespace experimental {
+
+// A wrapped fd that integrates with the EventEngine poller of the current
+// platform. A GrpcPolledFd knows how to create grpc platform-specific poller
+// handle from "ares_socket_t" sockets, and then sign up for
+// readability/writeability with that poller handle, and do shutdown and
+// destruction.
+class GrpcPolledFd {
+ public:
+  virtual ~GrpcPolledFd() {}
+  // Called when c-ares library is interested and there's no pending callback
+  virtual void RegisterForOnReadableLocked(
+      absl::AnyInvocable<void(absl::Status)> read_closure) = 0;
+  // Called when c-ares library is interested and there's no pending callback
+  virtual void RegisterForOnWriteableLocked(
+      absl::AnyInvocable<void(absl::Status)> write_closure) = 0;
+  // Indicates if there is data left even after just being read from
+  virtual bool IsFdStillReadableLocked() = 0;
+  // Called once and only once. Must cause cancellation of any pending
+  // read/write callbacks.
+  virtual void ShutdownLocked(grpc_error_handle error) = 0;
+  // Get the underlying ares_socket_t that this was created from
+  virtual ares_socket_t GetWrappedAresSocketLocked() = 0;
+  // A unique name, for logging
+  virtual const char* GetName() const = 0;
+};
+
+// A GrpcPolledFdFactory is 1-to-1 with and owned by a GrpcAresRequest. It knows
+// how to create GrpcPolledFd's for the current platform, and the
+// GrpcAresRequest uses it for all of its fd's.
+class GrpcPolledFdFactory {
+ public:
+  virtual ~GrpcPolledFdFactory() {}
+  // Creates a new wrapped fd for the current platform
+  virtual GrpcPolledFd* NewGrpcPolledFdLocked(ares_socket_t as) = 0;
+  // Optionally configures the ares channel after creation
+  virtual void ConfigureAresChannelLocked(ares_channel channel) = 0;
+};
+
+}  // namespace experimental
+}  // namespace grpc_event_engine
+
+#endif  // GRPC_ARES == 1
+#endif  // GRPC_SRC_CORE_LIB_EVENT_ENGINE_GRPC_POLLED_FD_H

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

@@ -316,7 +316,7 @@ void ResetEventManagerOnFork() {
   while (!fork_poller_list.empty()) {
     PollPoller* poller = fork_poller_list.front();
     fork_poller_list.pop_front();
-    delete poller;
+    poller->Close();
   }
   gpr_mu_unlock(&fork_fd_list_mu);
   InitPollPollerPosix();
@@ -566,6 +566,9 @@ bool PollEventHandle::EndPollLocked(bool got_read, bool got_write) {
 
 void PollPoller::KickExternal(bool ext) {
   grpc_core::MutexLock lock(&mu_);
+  if (closed_) {
+    return;
+  }
   if (was_kicked_) {
     if (ext) {
       was_kicked_ext_ = true;
@@ -610,7 +613,8 @@ PollPoller::PollPoller(Scheduler* scheduler)
       was_kicked_(false),
       was_kicked_ext_(false),
       num_poll_handles_(0),
-      poll_handles_list_head_(nullptr) {
+      poll_handles_list_head_(nullptr),
+      closed_(false) {
   wakeup_fd_ = *CreateWakeupFd();
   GPR_ASSERT(wakeup_fd_ != nullptr);
   ForkPollerListAddPoller(this);
@@ -622,7 +626,8 @@ PollPoller::PollPoller(Scheduler* scheduler, bool use_phony_poll)
       was_kicked_(false),
       was_kicked_ext_(false),
       num_poll_handles_(0),
-      poll_handles_list_head_(nullptr) {
+      poll_handles_list_head_(nullptr),
+      closed_(false) {
   wakeup_fd_ = *CreateWakeupFd();
   GPR_ASSERT(wakeup_fd_ != nullptr);
   ForkPollerListAddPoller(this);
@@ -829,6 +834,17 @@ void PollPoller::Shutdown() {
   Unref();
 }
 
+void PollPoller::PrepareFork() { Kick(); }
+// TODO(vigneshbabu): implement
+void PollPoller::PostforkParent() {}
+// TODO(vigneshbabu): implement
+void PollPoller::PostforkChild() {}
+
+void PollPoller::Close() {
+  grpc_core::MutexLock lock(&mu_);
+  closed_ = true;
+}
+
 PollPoller* MakePollPoller(Scheduler* scheduler, bool use_phony_poll) {
   static bool kPollPollerSupported = InitPollPollerPosix();
   if (kPollPollerSupported) {
@@ -875,6 +891,12 @@ PollPoller* MakePollPoller(Scheduler* /*scheduler*/,
   return nullptr;
 }
 
+void PollPoller::PrepareFork() { grpc_core::Crash("unimplemented"); }
+void PollPoller::PostforkParent() { grpc_core::Crash("unimplemented"); }
+void PollPoller::PostforkChild() { grpc_core::Crash("unimplemented"); }
+
+void PollPoller::Close() { grpc_core::Crash("unimplemented"); }
+
 void PollPoller::KickExternal(bool /*ext*/) {
   grpc_core::Crash("unimplemented");
 }

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

@@ -27,6 +27,7 @@
 
 #include <grpc/event_engine/event_engine.h>
 
+#include "src/core/lib/event_engine/forkable.h"
 #include "src/core/lib/event_engine/poller.h"
 #include "src/core/lib/event_engine/posix_engine/event_poller.h"
 #include "src/core/lib/event_engine/posix_engine/wakeup_fd_posix.h"
@@ -38,7 +39,7 @@ namespace experimental {
 class PollEventHandle;
 
 // Definition of poll based poller.
-class PollPoller : public PosixEventPoller {
+class PollPoller : public PosixEventPoller, public Forkable {
  public:
   explicit PollPoller(Scheduler* scheduler);
   PollPoller(Scheduler* scheduler, bool use_phony_poll);
@@ -54,6 +55,13 @@ class PollPoller : public PosixEventPoller {
   bool CanTrackErrors() const override { return false; }
   ~PollPoller() override;
 
+  // Forkable
+  void PrepareFork() override;
+  void PostforkParent() override;
+  void PostforkChild() override;
+
+  void Close();
+
  private:
   void Ref() { ref_count_.fetch_add(1, std::memory_order_relaxed); }
   void Unref() {
@@ -83,6 +91,7 @@ class PollPoller : public PosixEventPoller {
   int num_poll_handles_ ABSL_GUARDED_BY(mu_);
   PollEventHandle* poll_handles_list_head_ ABSL_GUARDED_BY(mu_) = nullptr;
   std::unique_ptr<WakeupFd> wakeup_fd_;
+  bool closed_ ABSL_GUARDED_BY(mu_);
 };
 
 // Return an instance of a poll based poller tied to the specified scheduler.

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

@@ -0,0 +1,197 @@
+// 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_POSIX_ENGINE_GRPC_POLLED_FD_POSIX_H
+#define GRPC_SRC_CORE_LIB_EVENT_ENGINE_POSIX_ENGINE_GRPC_POLLED_FD_POSIX_H
+
+#include <grpc/support/port_platform.h>
+
+#include "src/core/lib/iomgr/port.h"
+
+#if GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_ARES_EV_DRIVER)
+
+// IWYU pragma: no_include <ares_build.h>
+
+#include <sys/ioctl.h>
+#include <sys/socket.h>
+#include <sys/uio.h>
+#include <unistd.h>
+
+#include <string>
+#include <unordered_set>
+#include <utility>
+
+#include <ares.h>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+
+#include "src/core/lib/event_engine/grpc_polled_fd.h"
+#include "src/core/lib/event_engine/posix_engine/event_poller.h"
+#include "src/core/lib/event_engine/posix_engine/posix_engine_closure.h"
+#include "src/core/lib/event_engine/posix_engine/tcp_socket_utils.h"
+#include "src/core/lib/iomgr/error.h"
+
+namespace grpc_event_engine {
+namespace experimental {
+
+class GrpcPolledFdPosix : public GrpcPolledFd {
+ public:
+  GrpcPolledFdPosix(ares_socket_t as, EventHandle* handle)
+      : name_(absl::StrCat("c-ares fd: ", static_cast<int>(as))),
+        as_(as),
+        handle_(handle) {}
+
+  ~GrpcPolledFdPosix() override {
+    // c-ares library will close the fd. This fd may be picked up immediately by
+    // another thread and should not be closed by the following OrphanHandle.
+    int phony_release_fd;
+    handle_->OrphanHandle(/*on_done=*/nullptr, &phony_release_fd,
+                          "c-ares query finished");
+  }
+
+  void RegisterForOnReadableLocked(
+      absl::AnyInvocable<void(absl::Status)> read_closure) override {
+    handle_->NotifyOnRead(new PosixEngineClosure(std::move(read_closure),
+                                                 /*is_permanent=*/false));
+  }
+
+  void RegisterForOnWriteableLocked(
+      absl::AnyInvocable<void(absl::Status)> write_closure) override {
+    handle_->NotifyOnWrite(new PosixEngineClosure(std::move(write_closure),
+                                                  /*is_permanent=*/false));
+  }
+
+  bool IsFdStillReadableLocked() override {
+    size_t bytes_available = 0;
+    return ioctl(handle_->WrappedFd(), FIONREAD, &bytes_available) == 0 &&
+           bytes_available > 0;
+  }
+
+  void ShutdownLocked(grpc_error_handle error) override {
+    handle_->ShutdownHandle(error);
+  }
+
+  ares_socket_t GetWrappedAresSocketLocked() override { return as_; }
+
+  const char* GetName() const override { return name_.c_str(); }
+
+ private:
+  const std::string name_;
+  const ares_socket_t as_;
+  EventHandle* handle_;
+};
+
+class GrpcPolledFdFactoryPosix : public GrpcPolledFdFactory {
+ public:
+  explicit GrpcPolledFdFactoryPosix(PosixEventPoller* poller)
+      : poller_(poller) {}
+
+  ~GrpcPolledFdFactoryPosix() override {
+    for (auto& fd : owned_fds_) {
+      close(fd);
+    }
+  }
+
+  GrpcPolledFd* NewGrpcPolledFdLocked(ares_socket_t as) override {
+    owned_fds_.insert(as);
+    return new GrpcPolledFdPosix(
+        as,
+        poller_->CreateHandle(as, "c-ares socket", poller_->CanTrackErrors()));
+  }
+
+  void ConfigureAresChannelLocked(ares_channel channel) override {
+    ares_set_socket_functions(channel, &kSockFuncs, this);
+    ares_set_socket_configure_callback(
+        channel, &GrpcPolledFdFactoryPosix::ConfigureSocket, nullptr);
+  }
+
+ private:
+  /// Overridden socket API for c-ares
+  static ares_socket_t Socket(int af, int type, int protocol,
+                              void* /*user_data*/) {
+    return socket(af, type, protocol);
+  }
+
+  /// Overridden connect API for c-ares
+  static int Connect(ares_socket_t as, const struct sockaddr* target,
+                     ares_socklen_t target_len, void* /*user_data*/) {
+    return connect(as, target, target_len);
+  }
+
+  /// Overridden writev API for c-ares
+  static ares_ssize_t WriteV(ares_socket_t as, const struct iovec* iov,
+                             int iovec_count, void* /*user_data*/) {
+    return writev(as, iov, iovec_count);
+  }
+
+  /// Overridden recvfrom API for c-ares
+  static ares_ssize_t RecvFrom(ares_socket_t as, void* data, size_t data_len,
+                               int flags, struct sockaddr* from,
+                               ares_socklen_t* from_len, void* /*user_data*/) {
+    return recvfrom(as, data, data_len, flags, from, from_len);
+  }
+
+  /// Overridden close API for c-ares
+  static int Close(ares_socket_t as, void* user_data) {
+    GrpcPolledFdFactoryPosix* self =
+        static_cast<GrpcPolledFdFactoryPosix*>(user_data);
+    if (self->owned_fds_.find(as) == self->owned_fds_.end()) {
+      // c-ares owns this fd, grpc has never seen it
+      return close(as);
+    }
+    return 0;
+  }
+
+  /// Because we're using socket API overrides, c-ares won't
+  /// perform its typical configuration on the socket. See
+  /// https://github.com/c-ares/c-ares/blob/bad62225b7f6b278b92e8e85a255600b629ef517/src/lib/ares_process.c#L1018.
+  /// So we use the configure socket callback override and copy default
+  /// settings that c-ares would normally apply on posix platforms:
+  ///   - non-blocking
+  ///   - cloexec flag
+  ///   - disable nagle
+  static int ConfigureSocket(ares_socket_t fd, int type, void* /*user_data*/) {
+    // clang-format off
+#define RETURN_IF_ERROR(expr) if (!(expr).ok()) { return -1; }
+    // clang-format on
+    PosixSocketWrapper sock(fd);
+    RETURN_IF_ERROR(sock.SetSocketNonBlocking(1));
+    RETURN_IF_ERROR(sock.SetSocketCloexec(1));
+    if (type == SOCK_STREAM) {
+      RETURN_IF_ERROR(sock.SetSocketLowLatency(1));
+    }
+    return 0;
+  }
+
+  const struct ares_socket_functions kSockFuncs = {
+      &GrpcPolledFdFactoryPosix::Socket /* socket */,
+      &GrpcPolledFdFactoryPosix::Close /* close */,
+      &GrpcPolledFdFactoryPosix::Connect /* connect */,
+      &GrpcPolledFdFactoryPosix::RecvFrom /* recvfrom */,
+      &GrpcPolledFdFactoryPosix::WriteV /* writev */,
+  };
+
+  PosixEventPoller* poller_;
+  // fds that are used/owned by grpc - we (grpc) will close them rather than
+  // c-ares
+  std::unordered_set<ares_socket_t> owned_fds_;
+};
+
+}  // namespace experimental
+}  // namespace grpc_event_engine
+
+#endif  // GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_ARES_EV_DRIVER)
+#endif  // GRPC_SRC_CORE_LIB_EVENT_ENGINE_POSIX_ENGINE_GRPC_POLLED_FD_POSIX_H

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

@@ -521,7 +521,7 @@ void PosixEndpointImpl::UpdateRcvLowat() {
 void PosixEndpointImpl::MaybeMakeReadSlices() {
   static const int kBigAlloc = 64 * 1024;
   static const int kSmallAlloc = 8 * 1024;
-  if (incoming_buffer_->Length() < static_cast<size_t>(min_progress_size_)) {
+  if (incoming_buffer_->Length() < std::max<size_t>(min_progress_size_, 1)) {
     size_t allocate_length = min_progress_size_;
     const size_t target_length = static_cast<size_t>(target_length_);
     // If memory pressure is low and we think there will be more than
@@ -531,8 +531,8 @@ void PosixEndpointImpl::MaybeMakeReadSlices() {
     if (low_memory_pressure && target_length > allocate_length) {
       allocate_length = target_length;
     }
-    int extra_wanted =
-        allocate_length - static_cast<int>(incoming_buffer_->Length());
+    int extra_wanted = std::max<int>(
+        1, allocate_length - static_cast<int>(incoming_buffer_->Length()));
     if (extra_wanted >=
         (low_memory_pressure ? kSmallAlloc * 3 / 2 : kBigAlloc)) {
       while (extra_wanted > 0) {

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

@@ -18,6 +18,7 @@
 #include <algorithm>
 #include <atomic>
 #include <chrono>
+#include <cstdint>
 #include <cstring>
 #include <memory>
 #include <string>
@@ -37,8 +38,10 @@
 #include <grpc/support/log.h>
 
 #include "src/core/lib/debug/trace.h"
+#include "src/core/lib/event_engine/grpc_polled_fd.h"
 #include "src/core/lib/event_engine/poller.h"
 #include "src/core/lib/event_engine/posix.h"
+#include "src/core/lib/event_engine/posix_engine/grpc_polled_fd_posix.h"
 #include "src/core/lib/event_engine/posix_engine/tcp_socket_utils.h"
 #include "src/core/lib/event_engine/posix_engine/timer.h"
 #include "src/core/lib/event_engine/tcp_socket_utils.h"
@@ -472,6 +475,10 @@ void PosixEventEngine::Run(EventEngine::Closure* closure) {
 
 EventEngine::TaskHandle PosixEventEngine::RunAfterInternal(
     Duration when, absl::AnyInvocable<void()> cb) {
+  if (when <= Duration::zero()) {
+    Run(std::move(cb));
+    return TaskHandle::kInvalid;
+  }
   auto when_ts = ToTimestamp(timer_manager_.Now(), when);
   auto* cd = new ClosureData;
   cd->cb = std::move(cb);
@@ -487,10 +494,49 @@ EventEngine::TaskHandle PosixEventEngine::RunAfterInternal(
   return handle;
 }
 
+#if GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_TCP)
+
+PosixEventEngine::PosixDNSResolver::PosixDNSResolver(
+    grpc_core::OrphanablePtr<AresResolver> ares_resolver)
+    : ares_resolver_(std::move(ares_resolver)) {}
+
+void PosixEventEngine::PosixDNSResolver::LookupHostname(
+    LookupHostnameCallback on_resolve, absl::string_view name,
+    absl::string_view default_port) {
+  ares_resolver_->LookupHostname(name, default_port, std::move(on_resolve));
+}
+
+void PosixEventEngine::PosixDNSResolver::LookupSRV(LookupSRVCallback on_resolve,
+                                                   absl::string_view name) {
+  ares_resolver_->LookupSRV(name, std::move(on_resolve));
+}
+
+void PosixEventEngine::PosixDNSResolver::LookupTXT(LookupTXTCallback on_resolve,
+                                                   absl::string_view name) {
+  ares_resolver_->LookupTXT(name, std::move(on_resolve));
+}
+
+#endif  // GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_TCP)
+
 absl::StatusOr<std::unique_ptr<EventEngine::DNSResolver>>
 PosixEventEngine::GetDNSResolver(
-    EventEngine::DNSResolver::ResolverOptions const& /*options*/) {
+    const EventEngine::DNSResolver::ResolverOptions& options) {
+#if GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_TCP)
+  auto ares_resolver = AresResolver::CreateAresResolver(
+      options.dns_server,
+      std::make_unique<GrpcPolledFdFactoryPosix>(poller_manager_->Poller()),
+      shared_from_this());
+  if (!ares_resolver.ok()) {
+    return ares_resolver.status();
+  }
+  return std::make_unique<PosixEventEngine::PosixDNSResolver>(
+      std::move(*ares_resolver));
+#else   // GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_TCP)
+  // TODO(yijiem): Implement a basic A/AAAA-only native resolver in
+  // PosixEventEngine.
+  (void)options;
   grpc_core::Crash("unimplemented");
+#endif  // GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_TCP)
 }
 
 bool PosixEventEngine::IsWorkerThread() { grpc_core::Crash("unimplemented"); }

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

@@ -34,11 +34,13 @@
 #include <grpc/event_engine/event_engine.h>
 #include <grpc/event_engine/memory_allocator.h>
 
+#include "src/core/lib/event_engine/ares_resolver.h"
 #include "src/core/lib/event_engine/handle_containers.h"
 #include "src/core/lib/event_engine/posix.h"
 #include "src/core/lib/event_engine/posix_engine/event_poller.h"
 #include "src/core/lib/event_engine/posix_engine/timer_manager.h"
 #include "src/core/lib/event_engine/thread_pool/thread_pool.h"
+#include "src/core/lib/gprpp/orphanable.h"
 #include "src/core/lib/gprpp/sync.h"
 #include "src/core/lib/iomgr/port.h"
 #include "src/core/lib/surface/init_internally.h"
@@ -138,7 +140,11 @@ class PosixEventEngine final : public PosixEventEngineWithFdSupport,
  public:
   class PosixDNSResolver : public EventEngine::DNSResolver {
    public:
-    ~PosixDNSResolver() override;
+    PosixDNSResolver() = delete;
+#if GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_TCP)
+    explicit PosixDNSResolver(
+        grpc_core::OrphanablePtr<AresResolver> ares_resolver);
+#endif  // GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_TCP)
     void LookupHostname(LookupHostnameCallback on_resolve,
                         absl::string_view name,
                         absl::string_view default_port) override;
@@ -146,6 +152,11 @@ class PosixEventEngine final : public PosixEventEngineWithFdSupport,
                    absl::string_view name) override;
     void LookupTXT(LookupTXTCallback on_resolve,
                    absl::string_view name) override;
+
+#if GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_TCP)
+   private:
+    grpc_core::OrphanablePtr<AresResolver> ares_resolver_;
+#endif  // GRPC_ARES == 1 && defined(GRPC_POSIX_SOCKET_TCP)
   };
 
 #ifdef GRPC_POSIX_SOCKET_TCP

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

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/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