Track requests that could cause other requests to be created, and don't do

a real core shutdown of a CQ until such requests are done

2 files changed, 16 insertions, 4 deletions

diff --git a/src/cpp/common/completion_queue_cc.cc b/src/cpp/common/completion_queue_cc.cc
index 00cc102f92..558f47af84 100644
--- a/src/cpp/common/completion_queue_cc.cc
+++ b/src/cpp/common/completion_queue_cc.cc
@@ -43,11 +43,18 @@ namespace grpc {
 
 static internal::GrpcLibraryInitializer g_gli_initializer;
 
-CompletionQueue::CompletionQueue(grpc_completion_queue* take) : cq_(take) {}
+CompletionQueue::CompletionQueue(grpc_completion_queue* take) : cq_(take) {
+  RegisterAvalanching();
+}
+
+void CompletionQueue::Shutdown() { CompleteAvalanching(); }
 
-void CompletionQueue::Shutdown() {
-  g_gli_initializer.summon();
-  grpc_completion_queue_shutdown(cq_);
+void CompletionQueue::CompleteAvalanching() {
+  // Check if this was the last avalanching operation
+  if (gpr_atm_no_barrier_fetch_add(&avalanches_in_flight_,
+                                   static_cast<gpr_atm>(-1)) == 1) {
+    grpc_completion_queue_shutdown(cq_);
+  }
 }
 
 CompletionQueue::NextStatus CompletionQueue::AsyncNextInternal(
diff --git a/src/cpp/server/server_cc.cc b/src/cpp/server/server_cc.cc
index b7cfd6dbf1..96820ff963 100644
--- a/src/cpp/server/server_cc.cc
+++ b/src/cpp/server/server_cc.cc
@@ -575,9 +575,14 @@ ServerInterface::BaseAsyncRequest::BaseAsyncRequest(
       tag_(tag),
       delete_on_finalize_(delete_on_finalize),
       call_(nullptr) {
+  call_cq_->RegisterAvalanching(); // This op will trigger more ops
   memset(&initial_metadata_array_, 0, sizeof(initial_metadata_array_));
 }
 
+ServerInterface::BaseAsyncRequest::~BaseAsyncRequest() {
+  call_cq_->CompleteAvalanching();
+}
+
 bool ServerInterface::BaseAsyncRequest::FinalizeResult(void** tag,
                                                        bool* status) {
   if (*status) {