Refine shutdown API

1 files changed, 15 insertions, 8 deletions

diff --git a/src/core/surface/server.c b/src/core/surface/server.c
index 37ceca43ba..0bf7f8331f 100644
--- a/src/core/surface/server.c
+++ b/src/core/surface/server.c
@@ -124,6 +124,11 @@ struct channel_data {
   gpr_uint32 registered_method_max_probes;
 };
 
+typedef struct shutdown_tag {
+  void *tag;
+  grpc_completion_queue *cq;
+} shutdown_tag;
+
 struct grpc_server {
   size_t channel_filter_count;
   const grpc_channel_filter **channel_filters;
@@ -140,7 +145,7 @@ struct grpc_server {
 
   gpr_uint8 shutdown;
   size_t num_shutdown_tags;
-  void **shutdown_tags;
+  shutdown_tag *shutdown_tags;
 
   call_data *lists[CALL_LIST_COUNT];
   channel_data root_channel_data;
@@ -383,13 +388,11 @@ static int num_listeners(grpc_server *server) {
 }
 
 static void maybe_finish_shutdown(grpc_server *server) {
-  size_t i, j;
+  size_t i;
   if (server->shutdown && server->lists[ALL_CALLS] == NULL && server->listeners_destroyed == num_listeners(server)) {
     for (i = 0; i < server->num_shutdown_tags; i++) {
-      for (j = 0; j < server->cq_count; j++) {
-        grpc_cq_end_op(server->cqs[j], server->shutdown_tags[i],
-                       NULL, 1);
-      }
+      grpc_cq_end_op(server->shutdown_tags[i].cq, server->shutdown_tags[i].tag,
+                     NULL, 1);
     }
   }
 }
@@ -804,7 +807,8 @@ grpc_transport_setup_result grpc_server_setup_transport(
   return result;
 }
 
-void grpc_server_shutdown_and_notify(grpc_server *server, void *shutdown_tag) {
+void grpc_server_shutdown_and_notify(grpc_server *server,
+                                     grpc_completion_queue *cq, void *tag) {
   listener *l;
   requested_call_array requested_calls;
   channel_data **channels;
@@ -814,6 +818,7 @@ void grpc_server_shutdown_and_notify(grpc_server *server, void *shutdown_tag) {
   grpc_channel_op op;
   grpc_channel_element *elem;
   registered_method *rm;
+  shutdown_tag *sdt;
 
   /* lock, and gather up some stuff to do */
   gpr_mu_lock(&server->mu);
@@ -823,7 +828,9 @@ void grpc_server_shutdown_and_notify(grpc_server *server, void *shutdown_tag) {
   server->shutdown_tags =
       gpr_realloc(server->shutdown_tags,
                   sizeof(void *) * (server->num_shutdown_tags + 1));
-  server->shutdown_tags[server->num_shutdown_tags++] = shutdown_tag;
+  sdt = &server->shutdown_tags[server->num_shutdown_tags++];
+  sdt->tag = tag;
+  sdt->cq = cq;
   if (server->shutdown) {
     gpr_mu_unlock(&server->mu);
     return;