Merge pull request #6190 from dgquintas/lb_pollset_propagation

Enable use of pollset_set for a channel call

6 files changed, 63 insertions, 22 deletions

diff --git a/src/core/lib/surface/call.c b/src/core/lib/surface/call.c
index 74a09cf6ed..e899bc8098 100644
--- a/src/core/lib/surface/call.c
+++ b/src/core/lib/surface/call.c
@@ -110,6 +110,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;
@@ -218,13 +219,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;
@@ -261,9 +260,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");
@@ -318,10 +328,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
diff --git a/src/core/lib/surface/call.h b/src/core/lib/surface/call.h
index 2725e060b8..b640345c21 100644
--- a/src/core/lib/surface/call.h
+++ b/src/core/lib/surface/call.h
@@ -53,6 +53,8 @@ typedef void (*grpc_ioreq_completion_func)(grpc_exec_ctx *exec_ctx,
 grpc_call *grpc_call_create(grpc_channel *channel, grpc_call *parent_call,
                             uint32_t propagation_mask,
                             grpc_completion_queue *cq,
+                            /* if not NULL, it'll be used in lieu of \a cq */
+                            grpc_pollset_set *pollset_set_alternative,
                             const void *server_transport_data,
                             grpc_mdelem **add_initial_metadata,
                             size_t add_initial_metadata_count,
diff --git a/src/core/lib/surface/channel.c b/src/core/lib/surface/channel.c
index b6b760b5d8..a7ea6fa1f0 100644
--- a/src/core/lib/surface/channel.c
+++ b/src/core/lib/surface/channel.c
@@ -166,12 +166,14 @@ char *grpc_channel_get_target(grpc_channel *channel) {
 
 static grpc_call *grpc_channel_create_call_internal(
     grpc_channel *channel, grpc_call *parent_call, uint32_t propagation_mask,
-    grpc_completion_queue *cq, grpc_mdelem *path_mdelem,
-    grpc_mdelem *authority_mdelem, gpr_timespec deadline) {
+    grpc_completion_queue *cq, grpc_pollset_set *pollset_set_alternative,
+    grpc_mdelem *path_mdelem, grpc_mdelem *authority_mdelem,
+    gpr_timespec deadline) {
   grpc_mdelem *send_metadata[2];
   size_t num_metadata = 0;
 
   GPR_ASSERT(channel->is_client);
+  GPR_ASSERT(!(cq != NULL && pollset_set_alternative != NULL));
 
   send_metadata[num_metadata++] = path_mdelem;
   if (authority_mdelem != NULL) {
@@ -180,8 +182,9 @@ static grpc_call *grpc_channel_create_call_internal(
     send_metadata[num_metadata++] = GRPC_MDELEM_REF(channel->default_authority);
   }
 
-  return grpc_call_create(channel, parent_call, propagation_mask, cq, NULL,
-                          send_metadata, num_metadata, deadline);
+  return grpc_call_create(channel, parent_call, propagation_mask, cq,
+                          pollset_set_alternative, NULL, send_metadata,
+                          num_metadata, deadline);
 }
 
 grpc_call *grpc_channel_create_call(grpc_channel *channel,
@@ -201,7 +204,22 @@ grpc_call *grpc_channel_create_call(grpc_channel *channel,
            (int)deadline.clock_type, reserved));
   GPR_ASSERT(!reserved);
   return grpc_channel_create_call_internal(
-      channel, parent_call, propagation_mask, cq,
+      channel, parent_call, propagation_mask, cq, NULL,
+      grpc_mdelem_from_metadata_strings(GRPC_MDSTR_PATH,
+                                        grpc_mdstr_from_string(method)),
+      host ? grpc_mdelem_from_metadata_strings(GRPC_MDSTR_AUTHORITY,
+                                               grpc_mdstr_from_string(host))
+           : NULL,
+      deadline);
+}
+
+grpc_call *grpc_channel_create_pollset_set_call(
+    grpc_channel *channel, grpc_call *parent_call, uint32_t propagation_mask,
+    grpc_pollset_set *pollset_set, const char *method, const char *host,
+    gpr_timespec deadline, void *reserved) {
+  GPR_ASSERT(!reserved);
+  return grpc_channel_create_call_internal(
+      channel, parent_call, propagation_mask, NULL, pollset_set,
       grpc_mdelem_from_metadata_strings(GRPC_MDSTR_PATH,
                                         grpc_mdstr_from_string(method)),
       host ? grpc_mdelem_from_metadata_strings(GRPC_MDSTR_AUTHORITY,
@@ -245,7 +263,7 @@ grpc_call *grpc_channel_create_registered_call(
           (int)deadline.tv_nsec, (int)deadline.clock_type, reserved));
   GPR_ASSERT(!reserved);
   return grpc_channel_create_call_internal(
-      channel, parent_call, propagation_mask, completion_queue,
+      channel, parent_call, propagation_mask, completion_queue, NULL,
       GRPC_MDELEM_REF(rc->path),
       rc->authority ? GRPC_MDELEM_REF(rc->authority) : NULL, deadline);
 }
diff --git a/src/core/lib/surface/channel.h b/src/core/lib/surface/channel.h
index 22dae930e4..ff3debc31f 100644
--- a/src/core/lib/surface/channel.h
+++ b/src/core/lib/surface/channel.h
@@ -42,6 +42,11 @@ grpc_channel *grpc_channel_create(grpc_exec_ctx *exec_ctx, const char *target,
                                   grpc_channel_stack_type channel_stack_type,
                                   grpc_transport *optional_transport);
 
+grpc_call *grpc_channel_create_pollset_set_call(
+    grpc_channel *channel, grpc_call *parent_call, uint32_t propagation_mask,
+    grpc_pollset_set *pollset_set, const char *method, const char *host,
+    gpr_timespec deadline, void *reserved);
+
 /** Get a (borrowed) pointer to this channels underlying channel stack */
 grpc_channel_stack *grpc_channel_get_channel_stack(grpc_channel *channel);
 
diff --git a/src/core/lib/surface/lame_client.c b/src/core/lib/surface/lame_client.c
index eef862787f..8223555382 100644
--- a/src/core/lib/surface/lame_client.c
+++ b/src/core/lib/surface/lame_client.c
@@ -128,7 +128,7 @@ const grpc_channel_filter grpc_lame_filter = {
     lame_start_transport_op,
     sizeof(call_data),
     init_call_elem,
-    grpc_call_stack_ignore_set_pollset,
+    grpc_call_stack_ignore_set_pollset_or_pollset_set,
     destroy_call_elem,
     sizeof(channel_data),
     init_channel_elem,
diff --git a/src/core/lib/surface/server.c b/src/core/lib/surface/server.c
index 9532e090a4..88d898bc1c 100644
--- a/src/core/lib/surface/server.c
+++ b/src/core/lib/surface/server.c
@@ -793,9 +793,9 @@ static void accept_stream(grpc_exec_ctx *exec_ctx, void *cd,
                           const void *transport_server_data) {
   channel_data *chand = cd;
   /* create a call */
-  grpc_call *call =
-      grpc_call_create(chand->channel, NULL, 0, NULL, transport_server_data,
-                       NULL, 0, gpr_inf_future(GPR_CLOCK_MONOTONIC));
+  grpc_call *call = grpc_call_create(chand->channel, NULL, 0, NULL, NULL,
+                                     transport_server_data, NULL, 0,
+                                     gpr_inf_future(GPR_CLOCK_MONOTONIC));
   grpc_call_element *elem =
       grpc_call_stack_element(grpc_call_get_call_stack(call), 0);
   call_data *calld = elem->call_data;
@@ -910,7 +910,7 @@ const grpc_channel_filter grpc_server_top_filter = {
     grpc_channel_next_op,
     sizeof(call_data),
     init_call_elem,
-    grpc_call_stack_ignore_set_pollset,
+    grpc_call_stack_ignore_set_pollset_or_pollset_set,
     destroy_call_elem,
     sizeof(channel_data),
     init_channel_elem,