1 files changed, 22 insertions, 14 deletions

diff --git a/src/core/surface/completion_queue.c b/src/core/surface/completion_queue.c
index f0fa25712a..9e4c636bd5 100644
--- a/src/core/surface/completion_queue.c
+++ b/src/core/surface/completion_queue.c
@@ -80,7 +80,8 @@ grpc_completion_queue *grpc_completion_queue_create(void) {
   memset(cc, 0, sizeof(*cc));
   /* Initial ref is dropped by grpc_completion_queue_shutdown */
   gpr_ref_init(&cc->refs, 1);
-  gpr_ref_init(&cc->owning_refs, 1);
+  /* One for destroy(), one for pollset_shutdown */
+  gpr_ref_init(&cc->owning_refs, 2);
   grpc_pollset_init(&cc->pollset);
   return cc;
 }
@@ -91,14 +92,14 @@ void grpc_cq_internal_ref(grpc_completion_queue *cc) {
 
 static void on_pollset_destroy_done(void *arg) {
   grpc_completion_queue *cc = arg;
-  grpc_pollset_destroy(&cc->pollset);
-  gpr_free(cc);
+  grpc_cq_internal_unref(cc);
 }
 
 void grpc_cq_internal_unref(grpc_completion_queue *cc) {
   if (gpr_unref(&cc->owning_refs)) {
     GPR_ASSERT(cc->queue == NULL);
-    grpc_pollset_shutdown(&cc->pollset, on_pollset_destroy_done, cc);
+    grpc_pollset_destroy(&cc->pollset);
+    gpr_free(cc);
   }
 }
 
@@ -136,24 +137,24 @@ void grpc_cq_begin_op(grpc_completion_queue *cc, grpc_call *call) {
 
 /* Signal the end of an operation - if this is the last waiting-to-be-queued
    event, then enter shutdown mode */
-static void end_op_locked(grpc_completion_queue *cc,
-                          grpc_completion_type type) {
-  if (gpr_unref(&cc->refs)) {
-    GPR_ASSERT(!cc->shutdown);
-    GPR_ASSERT(cc->shutdown_called);
-    cc->shutdown = 1;
-  }
-}
-
 void grpc_cq_end_op(grpc_completion_queue *cc, void *tag, grpc_call *call,
                     int success) {
   event *ev;
+  int shutdown = 0;
   gpr_mu_lock(GRPC_POLLSET_MU(&cc->pollset));
   ev = add_locked(cc, GRPC_OP_COMPLETE, tag, call);
   ev->base.success = success;
-  end_op_locked(cc, GRPC_OP_COMPLETE);
+  if (gpr_unref(&cc->refs)) {
+    GPR_ASSERT(!cc->shutdown);
+    GPR_ASSERT(cc->shutdown_called);
+    cc->shutdown = 1;
+    shutdown = 1;
+  }
   gpr_mu_unlock(GRPC_POLLSET_MU(&cc->pollset));
   if (call) GRPC_CALL_INTERNAL_UNREF(call, "cq", 0);
+  if (shutdown) {
+    grpc_pollset_shutdown(&cc->pollset, on_pollset_destroy_done, cc);
+  }
 }
 
 /* Create a GRPC_QUEUE_SHUTDOWN event without queuing it anywhere */
@@ -169,6 +170,7 @@ grpc_event grpc_completion_queue_next(grpc_completion_queue *cc,
   event *ev = NULL;
   grpc_event ret;
 
+  grpc_cq_internal_ref(cc);
   gpr_mu_lock(GRPC_POLLSET_MU(&cc->pollset));
   for (;;) {
     if (cc->queue != NULL) {
@@ -200,6 +202,7 @@ grpc_event grpc_completion_queue_next(grpc_completion_queue *cc,
       memset(&ret, 0, sizeof(ret));
       ret.type = GRPC_QUEUE_TIMEOUT;
       GRPC_SURFACE_TRACE_RETURNED_EVENT(cc, &ret);
+      grpc_cq_internal_unref(cc);
       return ret;
     }
   }
@@ -207,6 +210,7 @@ grpc_event grpc_completion_queue_next(grpc_completion_queue *cc,
   ret = ev->base;
   gpr_free(ev);
   GRPC_SURFACE_TRACE_RETURNED_EVENT(cc, &ret);
+  grpc_cq_internal_unref(cc);
   return ret;
 }
 
@@ -244,6 +248,7 @@ grpc_event grpc_completion_queue_pluck(grpc_completion_queue *cc, void *tag,
   event *ev = NULL;
   grpc_event ret;
 
+  grpc_cq_internal_ref(cc);
   gpr_mu_lock(GRPC_POLLSET_MU(&cc->pollset));
   for (;;) {
     if ((ev = pluck_event(cc, tag))) {
@@ -258,6 +263,7 @@ grpc_event grpc_completion_queue_pluck(grpc_completion_queue *cc, void *tag,
       memset(&ret, 0, sizeof(ret));
       ret.type = GRPC_QUEUE_TIMEOUT;
       GRPC_SURFACE_TRACE_RETURNED_EVENT(cc, &ret);
+      grpc_cq_internal_unref(cc);
       return ret;
     }
   }
@@ -265,6 +271,7 @@ grpc_event grpc_completion_queue_pluck(grpc_completion_queue *cc, void *tag,
   ret = ev->base;
   gpr_free(ev);
   GRPC_SURFACE_TRACE_RETURNED_EVENT(cc, &ret);
+  grpc_cq_internal_unref(cc);
   return ret;
 }
 
@@ -280,6 +287,7 @@ void grpc_completion_queue_shutdown(grpc_completion_queue *cc) {
     GPR_ASSERT(!cc->shutdown);
     cc->shutdown = 1;
     gpr_mu_unlock(GRPC_POLLSET_MU(&cc->pollset));
+    grpc_pollset_shutdown(&cc->pollset, on_pollset_destroy_done, cc);
   }
 }