Add Server.Wait

1 files changed, 15 insertions, 10 deletions

diff --git a/src/cpp/server/server.cc b/src/cpp/server/server.cc
index f565d3aa5d..178fa1a716 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;