diff --git a/include/grpc++/server.h b/include/grpc++/server.h index 26d18d1bbe4be0c49600d1268eeb2c2d323e6646..429c0ff3cf8579b9f70b3c48ca87f85838780761 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 f565d3aa5d5ec16b1e936318585d4d8e842adf1c..178fa1a71679586c1a6ebefc435406179433a5a8 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;