Tweak completion queue to newly agreed upon rules

src/core/surface/completion_queue.c

@@ -83,7 +83,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);
   cc->allow_polling = 1;
   return cc;
@@ -95,14 +96,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);
   }
 }
 
@@ -145,25 +146,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;
-    gpr_cv_broadcast(GRPC_POLLSET_CV(&cc->pollset));
-  }
-}
-
 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;
+    gpr_cv_broadcast(GRPC_POLLSET_CV(&cc->pollset));
+    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 */
@@ -179,6 +179,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) {
@@ -214,6 +215,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;
     }
   }
@@ -221,6 +223,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;
 }
 
@@ -258,6 +261,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))) {
@@ -276,6 +280,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;
     }
   }
@@ -283,6 +288,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;
 }
 
@@ -299,6 +305,7 @@ void grpc_completion_queue_shutdown(grpc_completion_queue *cc) {
     cc->shutdown = 1;
     gpr_cv_broadcast(GRPC_POLLSET_CV(&cc->pollset));
     gpr_mu_unlock(GRPC_POLLSET_MU(&cc->pollset));
+    grpc_pollset_shutdown(&cc->pollset, on_pollset_destroy_done, cc);
   }
 }