From 05573f1692647ca30a10839b61731486700587ac Mon Sep 17 00:00:00 2001
From: yang-g <yangg@google.com>
Date: Mon, 11 Jul 2016 15:48:01 -0700
Subject: [PATCH] Make Server::Wait work for async only server.

---
 include/grpc++/server.h                |  2 ++
 src/cpp/server/server.cc               |  5 ++---
 test/cpp/end2end/async_end2end_test.cc | 18 ++++++++++++++++++
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/include/grpc++/server.h b/include/grpc++/server.h
index 7a8858ef194..9a2f8f11c42 100644
--- a/include/grpc++/server.h
+++ b/include/grpc++/server.h
@@ -183,6 +183,8 @@ class Server GRPC_FINAL : public ServerInterface, private GrpcLibraryCodegen {
   int num_running_cb_;
   grpc::condition_variable callback_cv_;
 
+  grpc::condition_variable shutdown_cv_;
+
   std::shared_ptr<GlobalCallbacks> global_callbacks_;
 
   std::list<SyncRequest>* sync_methods_;
diff --git a/src/cpp/server/server.cc b/src/cpp/server/server.cc
index fb4c68ebe49..374c9cbc048 100644
--- a/src/cpp/server/server.cc
+++ b/src/cpp/server/server.cc
@@ -462,14 +462,13 @@ void Server::ShutdownInternal(gpr_timespec deadline) {
     while (num_running_cb_ != 0) {
       callback_cv_.wait(lock);
     }
+    shutdown_cv_.notify_all();
   }
 }
 
 void Server::Wait() {
   grpc::unique_lock<grpc::mutex> lock(mu_);
-  while (num_running_cb_ != 0) {
-    callback_cv_.wait(lock);
-  }
+  shutdown_cv_.wait(lock);
 }
 
 void Server::PerformOpsOnCall(CallOpSetInterface* ops, Call* call) {
diff --git a/test/cpp/end2end/async_end2end_test.cc b/test/cpp/end2end/async_end2end_test.cc
index 6c7eae53a40..01ba8962d79 100644
--- a/test/cpp/end2end/async_end2end_test.cc
+++ b/test/cpp/end2end/async_end2end_test.cc
@@ -345,6 +345,24 @@ TEST_P(AsyncEnd2endTest, SequentialRpcs) {
   SendRpc(10);
 }
 
+// We do not need to protect notify because the use is synchronized.
+void ServerWait(Server* server, int* notify) {
+  server->Wait();
+  *notify = 1;
+}
+TEST_P(AsyncEnd2endTest, WaitAndShutdownTest) {
+  int notify = 0;
+  std::thread* wait_thread =
+      new std::thread(&ServerWait, server_.get(), &notify);
+  ResetStub();
+  SendRpc(1);
+  EXPECT_EQ(0, notify);
+  server_->Shutdown();
+  wait_thread->join();
+  EXPECT_EQ(1, notify);
+  delete wait_thread;
+}
+
 // Test a simple RPC using the async version of Next
 TEST_P(AsyncEnd2endTest, AsyncNextRpc) {
   ResetStub();