feat: gracefully fail adding to the CompletionQueue after Shutdown

google/cloud/grpc_utils/completion_queue.cc

@@ -91,7 +91,9 @@ CompletionQueue::MakeDeadlineTimer(
     std::chrono::system_clock::time_point deadline) {
   auto op = std::make_shared<AsyncTimerFuture>(impl_->CreateAlarm());
   void* tag = impl_->RegisterOperation(op);
-  op->Set(impl_->cq(), deadline, tag);
+  if (tag != nullptr) {
+    op->Set(impl_->cq(), deadline, tag);
+  }
   return op->GetFuture();
 }
 

google/cloud/grpc_utils/completion_queue_test.cc

@@ -284,6 +284,39 @@ TEST(CompletionQueueTest, RunAsync) {
   runner.join();
 }
 
+// Sets up a timer that reschedules itself and verfies we can shut down
+// cleanly whether we call `CancelAll()` on the queue first or not.
+namespace {
+void RunAndReschedule(CompletionQueue& cq, bool ok) {
+  if (ok) {
+    cq.MakeRelativeTimer(std::chrono::seconds(1))
+        .then([&cq](future<StatusOr<std::chrono::system_clock::time_point>>
+                        result) { RunAndReschedule(cq, result.get().ok()); });
+  }
+}
+}  // namespace
+
+TEST(CompletionQueueTest, ShutdownWithReschedulingTimer) {
+  CompletionQueue cq;
+  std::thread t([&cq] { cq.Run(); });
+
+  RunAndReschedule(cq, /*ok=*/true);
+
+  cq.Shutdown();
+  t.join();
+}
+
+TEST(CompletionQueueTest, CancelAndShutdownWithReschedulingTimer) {
+  CompletionQueue cq;
+  std::thread t([&cq] { cq.Run(); });
+
+  RunAndReschedule(cq, /*ok=*/true);
+
+  cq.CancelAll();
+  cq.Shutdown();
+  t.join();
+}
+
 }  // namespace
 }  // namespace GOOGLE_CLOUD_CPP_GRPC_UTILS_NS
 }  // namespace grpc_utils

google/cloud/grpc_utils/internal/async_read_stream_impl.h

@@ -158,7 +158,9 @@ class AsyncReadStreamImpl
     reader_ = async_call(context_.get(), request, &cq_->cq());
     auto callback = std::make_shared<NotifyStart>(this->shared_from_this());
     void* tag = cq_->RegisterOperation(std::move(callback));
-    reader_->StartCall(tag);
+    if (tag != nullptr) {
+      reader_->StartCall(tag);
+    }
   }
 
   /// Cancel the current streaming read RPC.
@@ -196,7 +198,9 @@ class AsyncReadStreamImpl
     auto callback = std::make_shared<NotifyRead>(this->shared_from_this());
     auto response = &callback->response;
     void* tag = cq_->RegisterOperation(std::move(callback));
-    reader_->Read(response, tag);
+    if (tag != nullptr) {
+      reader_->Read(response, tag);
+    }
   }
 
   /// Handle the result of a `Read()` call.
@@ -244,7 +248,9 @@ class AsyncReadStreamImpl
     auto callback = std::make_shared<NotifyFinish>(this->shared_from_this());
     auto status = &callback->status;
     void* tag = cq_->RegisterOperation(std::move(callback));
-    reader_->Finish(status, tag);
+    if (tag != nullptr) {
+      reader_->Finish(status, tag);
+    }
   }
 
   /// Handle the result of a Finish() request.
@@ -279,7 +285,9 @@ class AsyncReadStreamImpl
     auto callback = std::make_shared<NotifyDiscard>(this->shared_from_this());
     auto response = &callback->response;
     void* tag = cq_->RegisterOperation(std::move(callback));
-    reader_->Read(response, tag);
+    if (tag != nullptr) {
+      reader_->Read(response, tag);
+    }
   }
 
   /// Handle the result of a Discard() call.

google/cloud/grpc_utils/internal/completion_queue_impl.cc

@@ -78,6 +78,11 @@ void* CompletionQueueImpl::RegisterOperation(
     std::shared_ptr<AsyncGrpcOperation> op) {
   void* tag = op.get();
   std::unique_lock<std::mutex> lk(mu_);
+  if (shutdown_) {
+    lk.unlock();
+    op->Notify(/*ok=*/false);
+    return nullptr;
+  }
   auto ins =
       pending_ops_.emplace(reinterpret_cast<std::intptr_t>(tag), std::move(op));
   // After this point we no longer need the lock, so release it.

google/cloud/grpc_utils/internal/completion_queue_impl.h

@@ -54,8 +54,6 @@ class AsyncGrpcOperation : public AsyncOperation {
    * Derived classes wrap the callbacks provided by the application and invoke
    * the callback when this virtual member function is called.
    *
-   * @param cq the completion queue sending the notification, this is useful in
-   *   case the callback needs to retry the operation.
    * @param ok opaque parameter returned by grpc::CompletionQueue.  The
    *   semantics defined by gRPC depend on the type of operation, therefore the
    *   operation needs to interpret this parameter based on those semantics.