Merge pull request #1110 from yang-g/cancel

Set cq_ on servercontext

src/cpp/server/server.cc

@@ -107,6 +107,7 @@ class Server::SyncRequest GRPC_FINAL : public CompletionQueueTag {
           request_payload_(mrd->request_payload_),
           method_(mrd->method_) {
       ctx_.call_ = mrd->call_;
+      ctx_.cq_ = &cq_;
       GPR_ASSERT(mrd->in_flight_);
       mrd->in_flight_ = false;
       mrd->request_metadata_.count = 0;
@@ -364,6 +365,7 @@ class Server::AsyncRequest GRPC_FINAL : public CompletionQueueTag {
       }
     }
     ctx->call_ = call_;
+    ctx->cq_ = cq_;
     Call call(call_, server_, cq_);
     if (orig_status && call_) {
       ctx->BeginCompletionOp(&call);

test/cpp/end2end/end2end_test.cc

@@ -83,10 +83,30 @@ void MaybeEchoDeadline(ServerContext* context, const EchoRequest* request,
 
 class TestServiceImpl : public ::grpc::cpp::test::util::TestService::Service {
  public:
+  TestServiceImpl() : signal_client_(false) {}
+
   Status Echo(ServerContext* context, const EchoRequest* request,
               EchoResponse* response) GRPC_OVERRIDE {
     response->set_message(request->message());
     MaybeEchoDeadline(context, request, response);
+    if (request->has_param() && request->param().client_cancel_after_us()) {
+      {
+        std::unique_lock<std::mutex> lock(mu_);
+        signal_client_ = true;
+      }
+      while (!context->IsCancelled()) {
+        std::this_thread::sleep_for(std::chrono::microseconds(
+            request->param().client_cancel_after_us()));
+      }
+      return Status::Cancelled;
+    } else if (request->has_param() &&
+               request->param().server_cancel_after_us()) {
+      std::this_thread::sleep_for(
+          std::chrono::microseconds(request->param().server_cancel_after_us()));
+      return Status::Cancelled;
+    } else {
+      EXPECT_FALSE(context->IsCancelled());
+    }
     return Status::OK;
   }
 
@@ -130,6 +150,15 @@ class TestServiceImpl : public ::grpc::cpp::test::util::TestService::Service {
     }
     return Status::OK;
   }
+
+  bool signal_client() {
+    std::unique_lock<std::mutex> lock(mu_);
+    return signal_client_;
+  }
+
+ private:
+  bool signal_client_;
+  std::mutex mu_;
 };
 
 class TestServiceImplDupPkg
@@ -151,7 +180,8 @@ class End2endTest : public ::testing::Test {
     server_address_ << "localhost:" << port;
     // Setup server
     ServerBuilder builder;
-    builder.AddListeningPort(server_address_.str(), InsecureServerCredentials());
+    builder.AddListeningPort(server_address_.str(),
+                             InsecureServerCredentials());
     builder.RegisterService(&service_);
     builder.RegisterService(&dup_pkg_service_);
     builder.SetThreadPool(&thread_pool_);
@@ -423,6 +453,44 @@ TEST_F(End2endTest, BadCredentials) {
   EXPECT_EQ("Rpc sent on a lame channel.", s.details());
 }
 
+void CancelRpc(ClientContext* context, int delay_us, TestServiceImpl* service) {
+  std::this_thread::sleep_for(std::chrono::microseconds(delay_us));
+  while (!service->signal_client()) {
+  }
+  context->TryCancel();
+}
+
+// Client cancels rpc after 10ms
+TEST_F(End2endTest, ClientCancelsRpc) {
+  ResetStub();
+  EchoRequest request;
+  EchoResponse response;
+  request.set_message("Hello");
+  const int kCancelDelayUs = 10 * 1000;
+  request.mutable_param()->set_client_cancel_after_us(kCancelDelayUs);
+
+  ClientContext context;
+  std::thread cancel_thread(CancelRpc, &context, kCancelDelayUs, &service_);
+  Status s = stub_->Echo(&context, request, &response);
+  cancel_thread.join();
+  EXPECT_EQ(StatusCode::CANCELLED, s.code());
+  EXPECT_TRUE(s.details().empty());
+}
+
+// Server cancels rpc after 1ms
+TEST_F(End2endTest, ServerCancelsRpc) {
+  ResetStub();
+  EchoRequest request;
+  EchoResponse response;
+  request.set_message("Hello");
+  request.mutable_param()->set_server_cancel_after_us(1000);
+
+  ClientContext context;
+  Status s = stub_->Echo(&context, request, &response);
+  EXPECT_EQ(StatusCode::CANCELLED, s.code());
+  EXPECT_TRUE(s.details().empty());
+}
+
 }  // namespace testing
 }  // namespace grpc
 

test/cpp/util/messages.proto

@@ -34,6 +34,8 @@ package grpc.cpp.test.util;
 
 message RequestParams {
   optional bool echo_deadline = 1;
+  optional int32 client_cancel_after_us = 2;
+  optional int32 server_cancel_after_us = 3;
 }
 
 message EchoRequest {