1 files changed, 27 insertions, 11 deletions

diff --git a/src/core/lib/surface/call.c b/src/core/lib/surface/call.c
index 54c93f8829..a9b1e25a77 100644
--- a/src/core/lib/surface/call.c
+++ b/src/core/lib/surface/call.c
@@ -113,6 +113,7 @@ typedef struct batch_control {
 
 struct grpc_call {
   grpc_completion_queue *cq;
+  grpc_polling_entity pollent;
   grpc_channel *channel;
   grpc_call *parent;
   grpc_call *first_child;
@@ -224,13 +225,11 @@ static void destroy_call(grpc_exec_ctx *exec_ctx, void *call_stack,
 static void receiving_slice_ready(grpc_exec_ctx *exec_ctx, void *bctlp,
                                   bool success);
 
-grpc_call *grpc_call_create(grpc_channel *channel, grpc_call *parent_call,
-                            uint32_t propagation_mask,
-                            grpc_completion_queue *cq,
-                            const void *server_transport_data,
-                            grpc_mdelem **add_initial_metadata,
-                            size_t add_initial_metadata_count,
-                            gpr_timespec send_deadline) {
+grpc_call *grpc_call_create(
+    grpc_channel *channel, grpc_call *parent_call, uint32_t propagation_mask,
+    grpc_completion_queue *cq, grpc_pollset_set *pollset_set_alternative,
+    const void *server_transport_data, grpc_mdelem **add_initial_metadata,
+    size_t add_initial_metadata_count, gpr_timespec send_deadline) {
   size_t i, j;
   grpc_channel_stack *channel_stack = grpc_channel_get_channel_stack(channel);
   grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
@@ -267,9 +266,20 @@ grpc_call *grpc_call_create(grpc_channel *channel, grpc_call *parent_call,
                        call->context, server_transport_data,
                        CALL_STACK_FROM_CALL(call));
   if (cq != NULL) {
+    GPR_ASSERT(
+        pollset_set_alternative == NULL &&
+        "Only one of 'cq' and 'pollset_set_alternative' should be non-NULL.");
     GRPC_CQ_INTERNAL_REF(cq, "bind");
-    grpc_call_stack_set_pollset(&exec_ctx, CALL_STACK_FROM_CALL(call),
-                                grpc_cq_pollset(cq));
+    call->pollent =
+        grpc_polling_entity_create_from_pollset(grpc_cq_pollset(cq));
+  }
+  if (pollset_set_alternative != NULL) {
+    call->pollent =
+        grpc_polling_entity_create_from_pollset_set(pollset_set_alternative);
+  }
+  if (!grpc_polling_entity_is_empty(&call->pollent)) {
+    grpc_call_stack_set_pollset_or_pollset_set(
+        &exec_ctx, CALL_STACK_FROM_CALL(call), &call->pollent);
   }
   if (parent_call != NULL) {
     GRPC_CALL_INTERNAL_REF(parent_call, "child");
@@ -324,10 +334,16 @@ grpc_call *grpc_call_create(grpc_channel *channel, grpc_call *parent_call,
 void grpc_call_set_completion_queue(grpc_exec_ctx *exec_ctx, grpc_call *call,
                                     grpc_completion_queue *cq) {
   GPR_ASSERT(cq);
+
+  if (grpc_polling_entity_pollset_set(&call->pollent) != NULL) {
+    gpr_log(GPR_ERROR, "A pollset_set is already registered for this call.");
+    abort();
+  }
   call->cq = cq;
   GRPC_CQ_INTERNAL_REF(cq, "bind");
-  grpc_call_stack_set_pollset(exec_ctx, CALL_STACK_FROM_CALL(call),
-                              grpc_cq_pollset(cq));
+  call->pollent = grpc_polling_entity_create_from_pollset(grpc_cq_pollset(cq));
+  grpc_call_stack_set_pollset_or_pollset_set(
+      exec_ctx, CALL_STACK_FROM_CALL(call), &call->pollent);
 }
 
 #ifdef GRPC_STREAM_REFCOUNT_DEBUG