From 6e57b9edb19eaa302586beadcaa7cf21a3601025 Mon Sep 17 00:00:00 2001
From: Craig Tiller <ctiller@google.com>
Date: Tue, 24 Feb 2015 15:46:22 -0800
Subject: [PATCH] Add Server.Wait

---
 include/grpc++/server.h  |  5 +++++
 src/cpp/server/server.cc | 25 +++++++++++++++----------
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/include/grpc++/server.h b/include/grpc++/server.h
index 26d18d1bbe4..429c0ff3cf8 100644
--- a/include/grpc++/server.h
+++ b/include/grpc++/server.h
@@ -69,6 +69,11 @@ class Server final : private CallHook,
   // Shutdown the server, block until all rpc processing finishes.
   void Shutdown();
 
+  // Block waiting for all work to complete (the server must either
+  // be shutting down or some other thread must call Shutdown for this
+  // function to ever return)
+  void Wait();
+
  private:
   friend class ServerBuilder;
 
diff --git a/src/cpp/server/server.cc b/src/cpp/server/server.cc
index f565d3aa5d5..178fa1a7167 100644
--- a/src/cpp/server/server.cc
+++ b/src/cpp/server/server.cc
@@ -265,21 +265,26 @@ bool Server::Start() {
 }
 
 void Server::Shutdown() {
-  {
-    std::unique_lock<std::mutex> lock(mu_);
-    if (started_ && !shutdown_) {
-      shutdown_ = true;
-      grpc_server_shutdown(server_);
-      cq_.Shutdown();
+  std::unique_lock<std::mutex> lock(mu_);
+  if (started_ && !shutdown_) {
+    shutdown_ = true;
+    grpc_server_shutdown(server_);
+    cq_.Shutdown();
 
-      // Wait for running callbacks to finish.
-      while (num_running_cb_ != 0) {
-        callback_cv_.wait(lock);
-      }
+    // Wait for running callbacks to finish.
+    while (num_running_cb_ != 0) {
+      callback_cv_.wait(lock);
     }
   }
 }
 
+void Server::Wait() {
+  std::unique_lock<std::mutex> lock(mu_);
+  while (num_running_cb_ != 0) {
+    callback_cv_.wait(lock);
+  }
+}
+
 void Server::PerformOpsOnCall(CallOpBuffer* buf, Call* call) {
   static const size_t MAX_OPS = 8;
   size_t nops = MAX_OPS;