grpc 1.41.1 → 1.44.0.pre2

data/src/core/lib/iomgr/tcp_server_utils_posix_common.cc

@@ -26,6 +26,7 @@
 #include <limits.h>
 #include <stdio.h>
 #include <string.h>
+#include <sys/socket.h>
 
 #include <string>
 

data/src/core/lib/iomgr/tcp_server_utils_posix_ifaddrs.cc

@@ -26,6 +26,7 @@
 #include <ifaddrs.h>
 #include <stddef.h>
 #include <string.h>
+#include <sys/socket.h>
 
 #include <string>
 

data/src/core/lib/iomgr/tcp_server_windows.cc

@@ -96,15 +96,13 @@ struct grpc_tcp_server {
   grpc_closure* shutdown_complete;
 
   grpc_channel_args* channel_args;
-  grpc_slice_allocator_factory* slice_allocator_factory;
 };
 
 /* Public function. Allocates the proper data structures to hold a
    grpc_tcp_server. */
-static grpc_error_handle tcp_server_create(
-    grpc_closure* shutdown_complete, const grpc_channel_args* args,
-    grpc_slice_allocator_factory* slice_allocator_factory,
-    grpc_tcp_server** server) {
+static grpc_error_handle tcp_server_create(grpc_closure* shutdown_complete,
+                                           const grpc_channel_args* args,
+                                           grpc_tcp_server** server) {
   grpc_tcp_server* s = (grpc_tcp_server*)gpr_malloc(sizeof(grpc_tcp_server));
   s->channel_args = grpc_channel_args_copy(args);
   gpr_ref_init(&s->refs, 1);
@@ -117,7 +115,6 @@ static grpc_error_handle tcp_server_create(
   s->shutdown_starting.head = NULL;
   s->shutdown_starting.tail = NULL;
   s->shutdown_complete = shutdown_complete;
-  s->slice_allocator_factory = slice_allocator_factory;
   *server = s;
   return GRPC_ERROR_NONE;
 }
@@ -168,7 +165,6 @@ static void tcp_server_shutdown_starting_add(grpc_tcp_server* s,
 static void tcp_server_destroy(grpc_tcp_server* s) {
   grpc_tcp_listener* sp;
   gpr_mu_lock(&s->mu);
-  grpc_slice_allocator_factory_destroy(s->slice_allocator_factory);
   /* First, shutdown all fd's. This will queue abortion calls for all
      of the pending accepts due to the normal operation mechanism. */
   if (s->active_ports == 0) {
@@ -230,11 +226,10 @@ static grpc_error_handle prepare_socket(SOCKET sock,
 failure:
   GPR_ASSERT(error != GRPC_ERROR_NONE);
   grpc_error_set_int(
-      grpc_error_set_str(
-          GRPC_ERROR_CREATE_REFERENCING_FROM_STATIC_STRING(
-              "Failed to prepare server socket", &error, 1),
-          GRPC_ERROR_STR_TARGET_ADDRESS,
-          grpc_slice_from_cpp_string(grpc_sockaddr_to_uri(addr))),
+      grpc_error_set_str(GRPC_ERROR_CREATE_REFERENCING_FROM_STATIC_STRING(
+                             "Failed to prepare server socket", &error, 1),
+                         GRPC_ERROR_STR_TARGET_ADDRESS,
+                         grpc_sockaddr_to_uri(addr)),
       GRPC_ERROR_INT_FD, (intptr_t)sock);
   GRPC_ERROR_UNREF(error);
   if (sock != INVALID_SOCKET) closesocket(sock);
@@ -359,11 +354,8 @@ static void on_accept(void* arg, grpc_error_handle error) {
         gpr_free(utf8_message);
       }
       std::string fd_name = absl::StrCat("tcp_server:", peer_name_string);
-      ep = grpc_tcp_create(
-          grpc_winsocket_create(sock, fd_name.c_str()),
-          sp->server->channel_args, peer_name_string.c_str(),
-          grpc_slice_allocator_factory_create_slice_allocator(
-              sp->server->slice_allocator_factory, peer_name_string));
+      ep = grpc_tcp_create(grpc_winsocket_create(sock, fd_name.c_str()),
+                           sp->server->channel_args, peer_name_string);
     } else {
       closesocket(sock);
     }

data/src/core/lib/iomgr/tcp_windows.cc

@@ -116,8 +116,6 @@ typedef struct grpc_tcp {
   grpc_slice_buffer* write_slices;
   grpc_slice_buffer* read_slices;
 
-  grpc_slice_allocator* slice_allocator;
-
   /* The IO Completion Port runs from another thread. We need some mechanism
      to protect ourselves when requesting a shutdown. */
   gpr_mu mu;
@@ -132,7 +130,6 @@ static void tcp_free(grpc_tcp* tcp) {
   grpc_winsocket_destroy(tcp->socket);
   gpr_mu_destroy(&tcp->mu);
   grpc_slice_buffer_destroy_internal(&tcp->last_read_buffer);
-  grpc_slice_allocator_destroy(tcp->slice_allocator);
   if (tcp->shutting_down) GRPC_ERROR_UNREF(tcp->shutdown_error);
   delete tcp;
 }
@@ -186,7 +183,7 @@ static void on_read(void* tcpp, grpc_error_handle error) {
     gpr_log(GPR_INFO, "TCP:%p on_read", tcp);
   }
 
-  GRPC_ERROR_REF(error);
+  (void)GRPC_ERROR_REF(error);
 
   if (error == GRPC_ERROR_NONE) {
     if (info->wsa_error != 0 && !tcp->shutting_down) {
@@ -322,7 +319,7 @@ static void on_write(void* tcpp, grpc_error_handle error) {
     gpr_log(GPR_INFO, "TCP:%p on_write", tcp);
   }
 
-  GRPC_ERROR_REF(error);
+  (void)GRPC_ERROR_REF(error);
 
   gpr_mu_lock(&tcp->mu);
   cb = tcp->write_cb;
@@ -502,8 +499,7 @@ static grpc_endpoint_vtable vtable = {win_read,
 
 grpc_endpoint* grpc_tcp_create(grpc_winsocket* socket,
                                grpc_channel_args* channel_args,
-                               const char* peer_string,
-                               grpc_slice_allocator* slice_allocator) {
+                               absl::string_view peer_string) {
   grpc_tcp* tcp = new grpc_tcp;
   memset(tcp, 0, sizeof(grpc_tcp));
   tcp->base.vtable = &vtable;
@@ -521,9 +517,8 @@ grpc_endpoint* grpc_tcp_create(grpc_winsocket* socket,
   } else {
     tcp->local_address = grpc_sockaddr_to_uri(&resolved_local_addr);
   }
-  tcp->peer_string = peer_string;
+  tcp->peer_string = std::string(peer_string);
   grpc_slice_buffer_init(&tcp->last_read_buffer);
-  tcp->slice_allocator = slice_allocator;
   return &tcp->base;
 }
 

data/src/core/lib/iomgr/tcp_windows.h

@@ -42,8 +42,7 @@
  */
 grpc_endpoint* grpc_tcp_create(grpc_winsocket* socket,
                                grpc_channel_args* channel_args,
-                               const char* peer_string,
-                               grpc_slice_allocator* slice_allocator);
+                               absl::string_view peer_string);
 
 grpc_error_handle grpc_tcp_prepare_socket(SOCKET sock);
 

data/src/core/lib/iomgr/timer_generic.cc

@@ -43,7 +43,7 @@
 
 #define ADD_DEADLINE_SCALE 0.33
 #define MIN_QUEUE_WINDOW_DURATION 0.01
-#define MAX_QUEUE_WINDOW_DURATION 1
+#define MAX_QUEUE_WINDOW_DURATION 1.0
 
 grpc_core::TraceFlag grpc_timer_trace(false, "timer");
 grpc_core::TraceFlag grpc_timer_check_trace(false, "timer_check");
@@ -105,7 +105,7 @@ static void destroy_timer_ht() {
 }
 
 static bool is_in_ht(grpc_timer* t) {
-  size_t i = GPR_HASH_POINTER(t, NUM_HASH_BUCKETS);
+  size_t i = grpc_core::HashPointer(t, NUM_HASH_BUCKETS);
 
   gpr_mu_lock(&g_hash_mu[i]);
   grpc_timer* p = g_timer_ht[i];
@@ -119,7 +119,7 @@ static bool is_in_ht(grpc_timer* t) {
 
 static void add_to_ht(grpc_timer* t) {
   GPR_ASSERT(!t->hash_table_next);
-  size_t i = GPR_HASH_POINTER(t, NUM_HASH_BUCKETS);
+  size_t i = grpc_core::HashPointer(t, NUM_HASH_BUCKETS);
 
   gpr_mu_lock(&g_hash_mu[i]);
   grpc_timer* p = g_timer_ht[i];
@@ -144,7 +144,7 @@ static void add_to_ht(grpc_timer* t) {
 }
 
 static void remove_from_ht(grpc_timer* t) {
-  size_t i = GPR_HASH_POINTER(t, NUM_HASH_BUCKETS);
+  size_t i = grpc_core::HashPointer(t, NUM_HASH_BUCKETS);
   bool removed = false;
 
   gpr_mu_lock(&g_hash_mu[i]);
@@ -248,7 +248,7 @@ static grpc_millis compute_min_deadline(timer_shard* shard) {
 static void timer_list_init() {
   uint32_t i;
 
-  g_num_shards = GPR_CLAMP(2 * gpr_cpu_num_cores(), 1, 32);
+  g_num_shards = grpc_core::Clamp(2 * gpr_cpu_num_cores(), 1u, 32u);
   g_shards =
       static_cast<timer_shard*>(gpr_zalloc(g_num_shards * sizeof(*g_shards)));
   g_shard_queue = static_cast<timer_shard**>(
@@ -337,7 +337,7 @@ void grpc_timer_init_unset(grpc_timer* timer) { timer->pending = false; }
 static void timer_init(grpc_timer* timer, grpc_millis deadline,
                        grpc_closure* closure) {
   int is_first_timer = 0;
-  timer_shard* shard = &g_shards[GPR_HASH_POINTER(timer, g_num_shards)];
+  timer_shard* shard = &g_shards[grpc_core::HashPointer(timer, g_num_shards)];
   timer->closure = closure;
   timer->deadline = deadline;
 
@@ -444,7 +444,7 @@ static void timer_cancel(grpc_timer* timer) {
     return;
   }
 
-  timer_shard* shard = &g_shards[GPR_HASH_POINTER(timer, g_num_shards)];
+  timer_shard* shard = &g_shards[grpc_core::HashPointer(timer, g_num_shards)];
   gpr_mu_lock(&shard->mu);
   if (GRPC_TRACE_FLAG_ENABLED(grpc_timer_trace)) {
     gpr_log(GPR_INFO, "TIMER %p: CANCEL pending=%s", timer,
@@ -479,13 +479,13 @@ static bool refill_heap(timer_shard* shard, grpc_millis now) {
       grpc_time_averaged_stats_update_average(&shard->stats) *
       ADD_DEADLINE_SCALE;
   double deadline_delta =
-      GPR_CLAMP(computed_deadline_delta, MIN_QUEUE_WINDOW_DURATION,
-                MAX_QUEUE_WINDOW_DURATION);
+      grpc_core::Clamp(computed_deadline_delta, MIN_QUEUE_WINDOW_DURATION,
+                       MAX_QUEUE_WINDOW_DURATION);
   grpc_timer *timer, *next;
 
   /* Compute the new cap and put all timers under it into the queue: */
   shard->queue_deadline_cap =
-      saturating_add(GPR_MAX(now, shard->queue_deadline_cap),
+      saturating_add(std::max(now, shard->queue_deadline_cap),
                      static_cast<grpc_millis>(deadline_delta * 1000.0));
 
   if (GRPC_TRACE_FLAG_ENABLED(grpc_timer_check_trace)) {
@@ -583,7 +583,7 @@ static grpc_timer_check_result run_some_expired_timers(
   g_last_seen_min_timer = min_timer;
 
   if (now < min_timer) {
-    if (next != nullptr) *next = GPR_MIN(*next, min_timer);
+    if (next != nullptr) *next = std::min(*next, min_timer);
     return GRPC_TIMERS_CHECKED_AND_EMPTY;
   }
 
@@ -628,7 +628,7 @@ static grpc_timer_check_result run_some_expired_timers(
     }
 
     if (next) {
-      *next = GPR_MIN(*next, g_shard_queue[0]->min_deadline);
+      *next = std::min(*next, g_shard_queue[0]->min_deadline);
     }
 
 #if GPR_ARCH_64
@@ -663,7 +663,7 @@ static grpc_timer_check_result timer_check(grpc_millis* next) {
 
   if (now < min_timer) {
     if (next != nullptr) {
-      *next = GPR_MIN(*next, min_timer);
+      *next = std::min(*next, min_timer);
     }
     if (GRPC_TRACE_FLAG_ENABLED(grpc_timer_check_trace)) {
       gpr_log(GPR_INFO, "TIMER CHECK SKIP: now=%" PRId64 " min_timer=%" PRId64,

data/src/core/lib/iomgr/timer_heap.cc

@@ -97,7 +97,7 @@ void grpc_timer_heap_destroy(grpc_timer_heap* heap) { gpr_free(heap->timers); }
 bool grpc_timer_heap_add(grpc_timer_heap* heap, grpc_timer* timer) {
   if (heap->timer_count == heap->timer_capacity) {
     heap->timer_capacity =
-        GPR_MAX(heap->timer_capacity + 1, heap->timer_capacity * 3 / 2);
+        std::max(heap->timer_capacity + 1, heap->timer_capacity * 3 / 2);
     heap->timers = static_cast<grpc_timer**>(
         gpr_realloc(heap->timers, heap->timer_capacity * sizeof(grpc_timer*)));
   }

data/src/core/lib/iomgr/unix_sockets_posix.cc

@@ -35,29 +35,35 @@
 #include "src/core/lib/gpr/useful.h"
 #include "src/core/lib/iomgr/sockaddr.h"
 #include "src/core/lib/iomgr/unix_sockets_posix.h"
+#include "src/core/lib/transport/error_utils.h"
 
 void grpc_create_socketpair_if_unix(int sv[2]) {
   GPR_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv) == 0);
 }
 
-grpc_error_handle grpc_resolve_unix_domain_address(
-    const char* name, grpc_resolved_addresses** addresses) {
-  *addresses = static_cast<grpc_resolved_addresses*>(
-      gpr_malloc(sizeof(grpc_resolved_addresses)));
-  (*addresses)->naddrs = 1;
-  (*addresses)->addrs = static_cast<grpc_resolved_address*>(
-      gpr_malloc(sizeof(grpc_resolved_address)));
-  return grpc_core::UnixSockaddrPopulate(name, (*addresses)->addrs);
+absl::StatusOr<std::vector<grpc_resolved_address>>
+grpc_resolve_unix_domain_address(absl::string_view name) {
+  grpc_resolved_address addr;
+  grpc_error_handle error = grpc_core::UnixSockaddrPopulate(name, &addr);
+  if (error == GRPC_ERROR_NONE) {
+    return std::vector<grpc_resolved_address>({addr});
+  }
+  auto result = grpc_error_to_absl_status(error);
+  GRPC_ERROR_UNREF(error);
+  return result;
 }
 
-grpc_error_handle grpc_resolve_unix_abstract_domain_address(
-    const absl::string_view name, grpc_resolved_addresses** addresses) {
-  *addresses = static_cast<grpc_resolved_addresses*>(
-      gpr_malloc(sizeof(grpc_resolved_addresses)));
-  (*addresses)->naddrs = 1;
-  (*addresses)->addrs = static_cast<grpc_resolved_address*>(
-      gpr_malloc(sizeof(grpc_resolved_address)));
-  return grpc_core::UnixAbstractSockaddrPopulate(name, (*addresses)->addrs);
+absl::StatusOr<std::vector<grpc_resolved_address>>
+grpc_resolve_unix_abstract_domain_address(const absl::string_view name) {
+  grpc_resolved_address addr;
+  grpc_error_handle error =
+      grpc_core::UnixAbstractSockaddrPopulate(name, &addr);
+  if (error == GRPC_ERROR_NONE) {
+    return std::vector<grpc_resolved_address>({addr});
+  }
+  auto result = grpc_error_to_absl_status(error);
+  GRPC_ERROR_UNREF(error);
+  return result;
 }
 
 int grpc_is_unix_socket(const grpc_resolved_address* resolved_addr) {
@@ -87,22 +93,4 @@ void grpc_unlink_if_unix_domain_socket(
   }
 }
 
-std::string grpc_sockaddr_to_uri_unix_if_possible(
-    const grpc_resolved_address* resolved_addr) {
-  const grpc_sockaddr* addr =
-      reinterpret_cast<const grpc_sockaddr*>(resolved_addr->addr);
-  if (addr->sa_family != AF_UNIX) {
-    return "";
-  }
-  const auto* unix_addr = reinterpret_cast<const struct sockaddr_un*>(addr);
-  if (unix_addr->sun_path[0] == '\0' && unix_addr->sun_path[1] != '\0') {
-    return absl::StrCat(
-        "unix-abstract:",
-        absl::string_view(
-            unix_addr->sun_path + 1,
-            resolved_addr->len - sizeof(unix_addr->sun_family) - 1));
-  }
-  return absl::StrCat("unix:", unix_addr->sun_path);
-}
-
 #endif

data/src/core/lib/iomgr/unix_sockets_posix.h

@@ -32,18 +32,15 @@
 
 void grpc_create_socketpair_if_unix(int sv[2]);
 
-grpc_error_handle grpc_resolve_unix_domain_address(
-    const char* name, grpc_resolved_addresses** addresses);
+absl::StatusOr<std::vector<grpc_resolved_address>>
+grpc_resolve_unix_domain_address(absl::string_view name);
 
-grpc_error_handle grpc_resolve_unix_abstract_domain_address(
-    absl::string_view name, grpc_resolved_addresses** addresses);
+absl::StatusOr<std::vector<grpc_resolved_address>>
+grpc_resolve_unix_abstract_domain_address(absl::string_view name);
 
 int grpc_is_unix_socket(const grpc_resolved_address* resolved_addr);
 
 void grpc_unlink_if_unix_domain_socket(
     const grpc_resolved_address* resolved_addr);
 
-std::string grpc_sockaddr_to_uri_unix_if_possible(
-    const grpc_resolved_address* resolved_addr);
-
 #endif /* GRPC_CORE_LIB_IOMGR_UNIX_SOCKETS_POSIX_H */

data/src/core/lib/iomgr/unix_sockets_posix_noop.cc

@@ -33,18 +33,14 @@ void grpc_create_socketpair_if_unix(int /* sv */[2]) {
   GPR_ASSERT(0);
 }
 
-grpc_error_handle grpc_resolve_unix_domain_address(
-    const char* /* name */, grpc_resolved_addresses** addresses) {
-  *addresses = NULL;
-  return GRPC_ERROR_CREATE_FROM_STATIC_STRING(
-      "Unix domain sockets are not supported on Windows");
+absl::StatusOr<std::vector<grpc_resolved_address>>
+grpc_resolve_unix_domain_address(absl::string_view /* name */) {
+  return absl::UnknownError("Unix domain sockets are not supported on Windows");
 }
 
-grpc_error_handle grpc_resolve_unix_abstract_domain_address(
-    absl::string_view, grpc_resolved_addresses** addresses) {
-  *addresses = NULL;
-  return GRPC_ERROR_CREATE_FROM_STATIC_STRING(
-      "Unix domain sockets are not supported on Windows");
+absl::StatusOr<std::vector<grpc_resolved_address>>
+grpc_resolve_unix_abstract_domain_address(absl::string_view /* name */) {
+  return absl::UnknownError("Unix domain sockets are not supported on Windows");
 }
 
 int grpc_is_unix_socket(const grpc_resolved_address* /* addr */) {
@@ -54,9 +50,4 @@ int grpc_is_unix_socket(const grpc_resolved_address* /* addr */) {
 void grpc_unlink_if_unix_domain_socket(
     const grpc_resolved_address* /* addr */) {}
 
-std::string grpc_sockaddr_to_uri_unix_if_possible(
-    const grpc_resolved_address* /* addr */) {
-  return "";
-}
-
 #endif

data/src/core/lib/iomgr/work_serializer.cc

@@ -24,54 +24,79 @@ namespace grpc_core {
 
 DebugOnlyTraceFlag grpc_work_serializer_trace(false, "work_serializer");
 
-struct CallbackWrapper {
-  CallbackWrapper(std::function<void()> cb, const grpc_core::DebugLocation& loc)
-      : callback(std::move(cb)), location(loc) {}
-
-  MultiProducerSingleConsumerQueue::Node mpscq_node;
-  const std::function<void()> callback;
-  const DebugLocation location;
-};
-
 class WorkSerializer::WorkSerializerImpl : public Orphanable {
  public:
-  void Run(std::function<void()> callback,
-           const grpc_core::DebugLocation& location);
-
+  void Run(std::function<void()> callback, const DebugLocation& location);
+  void Schedule(std::function<void()> callback, const DebugLocation& location);
+  void DrainQueue();
   void Orphan() override;
 
  private:
-  void DrainQueue();
+  struct CallbackWrapper {
+    CallbackWrapper(std::function<void()> cb, const DebugLocation& loc)
+        : callback(std::move(cb)), location(loc) {}
+
+    MultiProducerSingleConsumerQueue::Node mpscq_node;
+    const std::function<void()> callback;
+    const DebugLocation location;
+  };
+
+  // Callers of DrainQueueOwned should make sure to grab the lock on the
+  // workserializer with
+  //
+  //   prev_ref_pair =
+  //     refs_.fetch_add(MakeRefPair(1, 1), std::memory_order_acq_rel);
+  //
+  // and only invoke DrainQueueOwned() if there was previously no owner. Note
+  // that the queue size is also incremented as part of the fetch_add to allow
+  // the callers to add a callback to the queue if another thread already holds
+  // the lock to the work serializer.
+  void DrainQueueOwned();
+
+  // First 16 bits indicate ownership of the WorkSerializer, next 48 bits are
+  // queue size (i.e., refs).
+  static uint64_t MakeRefPair(uint16_t owners, uint64_t size) {
+    GPR_ASSERT(size >> 48 == 0);
+    return (static_cast<uint64_t>(owners) << 48) + static_cast<int64_t>(size);
+  }
+  static uint32_t GetOwners(uint64_t ref_pair) {
+    return static_cast<uint32_t>(ref_pair >> 48);
+  }
+  static uint32_t GetSize(uint64_t ref_pair) {
+    return static_cast<uint32_t>(ref_pair & 0xffffffffffffu);
+  }
 
   // An initial size of 1 keeps track of whether the work serializer has been
   // orphaned.
-  std::atomic<size_t> size_{1};
+  std::atomic<uint64_t> refs_{MakeRefPair(0, 1)};
   MultiProducerSingleConsumerQueue queue_;
 };
 
-void WorkSerializer::WorkSerializerImpl::Run(
-    std::function<void()> callback, const grpc_core::DebugLocation& location) {
+void WorkSerializer::WorkSerializerImpl::Run(std::function<void()> callback,
+                                             const DebugLocation& location) {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
     gpr_log(GPR_INFO, "WorkSerializer::Run() %p Scheduling callback [%s:%d]",
             this, location.file(), location.line());
   }
-  const size_t prev_size = size_.fetch_add(1);
+  // Increment queue size for the new callback and owner count to attempt to
+  // take ownership of the WorkSerializer.
+  const uint64_t prev_ref_pair =
+      refs_.fetch_add(MakeRefPair(1, 1), std::memory_order_acq_rel);
   // The work serializer should not have been orphaned.
-  GPR_DEBUG_ASSERT(prev_size > 0);
-  if (prev_size == 1) {
-    // There is no other closure executing right now on this work serializer.
-    // Execute this closure immediately.
+  GPR_DEBUG_ASSERT(GetSize(prev_ref_pair) > 0);
+  if (GetOwners(prev_ref_pair) == 0) {
+    // We took ownership of the WorkSerializer. Invoke callback and drain queue.
     if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
       gpr_log(GPR_INFO, "  Executing immediately");
     }
     callback();
-    // Loan this thread to the work serializer thread and drain the queue.
-    DrainQueue();
+    DrainQueueOwned();
   } else {
+    // Another thread is holding the WorkSerializer, so decrement the ownership
+    // count we just added and queue the callback.
+    refs_.fetch_sub(MakeRefPair(1, 0), std::memory_order_acq_rel);
     CallbackWrapper* cb_wrapper =
         new CallbackWrapper(std::move(callback), location);
-    // There already are closures executing on this work serializer. Simply add
-    // this closure to the queue.
     if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
       gpr_log(GPR_INFO, "  Scheduling on queue : item %p", cb_wrapper);
     }
@@ -79,12 +104,27 @@ void WorkSerializer::WorkSerializerImpl::Run(
   }
 }
 
+void WorkSerializer::WorkSerializerImpl::Schedule(
+    std::function<void()> callback, const DebugLocation& location) {
+  CallbackWrapper* cb_wrapper =
+      new CallbackWrapper(std::move(callback), location);
+  if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
+    gpr_log(GPR_INFO,
+            "WorkSerializer::Schedule() %p Scheduling callback %p [%s:%d]",
+            this, cb_wrapper, location.file(), location.line());
+  }
+  refs_.fetch_add(MakeRefPair(0, 1), std::memory_order_acq_rel);
+  queue_.Push(&cb_wrapper->mpscq_node);
+}
+
 void WorkSerializer::WorkSerializerImpl::Orphan() {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
     gpr_log(GPR_INFO, "WorkSerializer::Orphan() %p", this);
   }
-  size_t prev_size = size_.fetch_sub(1);
-  if (prev_size == 1) {
+  uint64_t prev_ref_pair =
+      refs_.fetch_sub(MakeRefPair(0, 1), std::memory_order_acq_rel);
+  if (GetSize(prev_ref_pair) == 1) {
+    GPR_DEBUG_ASSERT(GetOwners(prev_ref_pair) == 0);
     if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
       gpr_log(GPR_INFO, "  Destroying");
     }
@@ -93,30 +133,52 @@ void WorkSerializer::WorkSerializerImpl::Orphan() {
 }
 
 // The thread that calls this loans itself to the work serializer so as to
-// execute all the scheduled callback. This is called from within
-// WorkSerializer::Run() after executing a callback immediately, and hence size_
-// is at least 1.
+// execute all the scheduled callbacks.
 void WorkSerializer::WorkSerializerImpl::DrainQueue() {
+  if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
+    gpr_log(GPR_INFO, "WorkSerializer::DrainQueue() %p", this);
+  }
+  // Attempt to take ownership of the WorkSerializer. Also increment the queue
+  // size as required by `DrainQueueOwned()`.
+  const uint64_t prev_ref_pair =
+      refs_.fetch_add(MakeRefPair(1, 1), std::memory_order_acq_rel);
+  if (GetOwners(prev_ref_pair) == 0) {
+    // We took ownership of the WorkSerializer. Drain the queue.
+    DrainQueueOwned();
+  } else {
+    // Another thread is holding the WorkSerializer, so decrement the ownership
+    // count we just added and queue a no-op callback.
+    refs_.fetch_sub(MakeRefPair(1, 0), std::memory_order_acq_rel);
+    CallbackWrapper* cb_wrapper = new CallbackWrapper([]() {}, DEBUG_LOCATION);
+    queue_.Push(&cb_wrapper->mpscq_node);
+  }
+}
+
+void WorkSerializer::WorkSerializerImpl::DrainQueueOwned() {
+  if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
+    gpr_log(GPR_INFO, "WorkSerializer::DrainQueueOwned() %p", this);
+  }
   while (true) {
-    if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
-      gpr_log(GPR_INFO, "WorkSerializer::DrainQueue() %p", this);
-    }
-    size_t prev_size = size_.fetch_sub(1);
-    GPR_DEBUG_ASSERT(prev_size >= 1);
-    // It is possible that while draining the queue, one of the callbacks ended
+    auto prev_ref_pair = refs_.fetch_sub(MakeRefPair(0, 1));
+    // It is possible that while draining the queue, the last callback ended
     // up orphaning the work serializer. In that case, delete the object.
-    if (prev_size == 1) {
+    if (GetSize(prev_ref_pair) == 1) {
       if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
         gpr_log(GPR_INFO, "  Queue Drained. Destroying");
       }
       delete this;
       return;
     }
-    if (prev_size == 2) {
-      if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
-        gpr_log(GPR_INFO, "  Queue Drained");
+    if (GetSize(prev_ref_pair) == 2) {
+      // Queue drained. Give up ownership but only if queue remains empty. Note
+      // that we are using relaxed memory order semantics for the load on
+      // failure since we don't care about that value.
+      uint64_t expected = MakeRefPair(1, 1);
+      if (refs_.compare_exchange_strong(expected, MakeRefPair(0, 1),
+                                        std::memory_order_acq_rel,
+                                        std::memory_order_relaxed)) {
+        return;
       }
-      return;
     }
     // There is at least one callback on the queue. Pop the callback from the
     // queue and execute it.
@@ -124,8 +186,8 @@ void WorkSerializer::WorkSerializerImpl::DrainQueue() {
     bool empty_unused;
     while ((cb_wrapper = reinterpret_cast<CallbackWrapper*>(
                 queue_.PopAndCheckEnd(&empty_unused))) == nullptr) {
-      // This can happen either due to a race condition within the mpscq
-      // implementation or because of a race with Run()
+      // This can happen due to a race condition within the mpscq
+      // implementation or because of a race with Run()/Schedule().
       if (GRPC_TRACE_FLAG_ENABLED(grpc_work_serializer_trace)) {
         gpr_log(GPR_INFO, "  Queue returned nullptr, trying again");
       }
@@ -140,7 +202,9 @@ void WorkSerializer::WorkSerializerImpl::DrainQueue() {
   }
 }
 
+//
 // WorkSerializer
+//
 
 WorkSerializer::WorkSerializer()
     : impl_(MakeOrphanable<WorkSerializerImpl>()) {}
@@ -148,8 +212,15 @@ WorkSerializer::WorkSerializer()
 WorkSerializer::~WorkSerializer() {}
 
 void WorkSerializer::Run(std::function<void()> callback,
-                         const grpc_core::DebugLocation& location) {
+                         const DebugLocation& location) {
   impl_->Run(std::move(callback), location);
 }
 
+void WorkSerializer::Schedule(std::function<void()> callback,
+                              const DebugLocation& location) {
+  impl_->Schedule(std::move(callback), location);
+}
+
+void WorkSerializer::DrainQueue() { impl_->DrainQueue(); }
+
 }  // namespace grpc_core