Merge pull request grpc#761 from ctiller/chexxy

Add Server.Wait
rainbowsworld · Feb 24, 2015 · ab399fa · ab399fa
2 parents 175533c + 6e57b9e
commit ab399fa
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git 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
@@ -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;