Properly integrate async API with server-side cancellations.

There is a comment above IsCancelled that says when it is ok to use this.

1 files changed, 18 insertions, 8 deletions

diff --git a/src/cpp/server/server_context.cc b/src/cpp/server/server_context.cc
index e205a1969b..eb49b21037 100644
--- a/src/cpp/server/server_context.cc
+++ b/src/cpp/server/server_context.cc
@@ -62,7 +62,11 @@ class ServerContext::CompletionOp GRPC_FINAL : public CallOpSetInterface {
   void FillOps(grpc_op* ops, size_t* nops) GRPC_OVERRIDE;
   bool FinalizeResult(void** tag, bool* status) GRPC_OVERRIDE;
 
-  bool CheckCancelled(CompletionQueue* cq);
+  bool CheckCancelled(CompletionQueue* cq) {
+    cq->TryPluck(this);
+    return CheckCancelledNoPluck();
+  }
+  bool CheckCancelledAsync() { return CheckCancelledNoPluck(); }
 
   void set_tag(void* tag) {
     has_tag_ = true;
@@ -72,6 +76,11 @@ class ServerContext::CompletionOp GRPC_FINAL : public CallOpSetInterface {
   void Unref();
 
  private:
+  bool CheckCancelledNoPluck() {
+    grpc::lock_guard<grpc::mutex> g(mu_);
+    return finalized_ ? (cancelled_ != 0) : false;
+  }
+
   bool has_tag_;
   void* tag_;
   grpc::mutex mu_;
@@ -88,12 +97,6 @@ void ServerContext::CompletionOp::Unref() {
   }
 }
 
-bool ServerContext::CompletionOp::CheckCancelled(CompletionQueue* cq) {
-  cq->TryPluck(this);
-  grpc::lock_guard<grpc::mutex> g(mu_);
-  return finalized_ ? cancelled_ != 0 : false;
-}
-
 void ServerContext::CompletionOp::FillOps(grpc_op* ops, size_t* nops) {
   ops->op = GRPC_OP_RECV_CLOSE_ON_SERVER;
   ops->data.recv_close_on_server.cancelled = &cancelled_;
@@ -182,7 +185,14 @@ void ServerContext::TryCancel() const {
 }
 
 bool ServerContext::IsCancelled() const {
-  return completion_op_ && completion_op_->CheckCancelled(cq_);
+  if (has_notify_when_done_tag_) {
+    // when using async API, but the result is only valid
+    // if the tag has already been delivered at the completion queue
+    return completion_op_ && completion_op_->CheckCancelledAsync();
+  } else {
+    // when using sync API
+    return completion_op_ && completion_op_->CheckCancelled(cq_);
+  }
 }
 
 void ServerContext::set_compression_level(grpc_compression_level level) {