Minor changes to GrpcRpcManager

2 files changed, 9 insertions, 8 deletions

diff --git a/src/cpp/rpcmanager/grpc_rpc_manager.cc b/src/cpp/rpcmanager/grpc_rpc_manager.cc
index 5bfd4eedad..1c7d5adeaf 100644
--- a/src/cpp/rpcmanager/grpc_rpc_manager.cc
+++ b/src/cpp/rpcmanager/grpc_rpc_manager.cc
@@ -83,7 +83,7 @@ void GrpcRpcManager::Wait() {
 }
 
 // For testing only
-void GrpcRpcManager::Shutdown() {
+void GrpcRpcManager::ShutdownRpcManager() {
   std::unique_lock<grpc::mutex> lock(mu_);
   shutdown_ = true;
 }
@@ -131,9 +131,10 @@ void GrpcRpcManager::MaybeCreatePoller() {
 
 void GrpcRpcManager::MainWorkLoop() {
   bool is_work_found = false;
+  void *tag;
 
   do {
-    PollForWork(is_work_found);
+    PollForWork(is_work_found, &tag);
 
     // Decrement num_pollers since this thread is no longer polling
     {
@@ -146,7 +147,7 @@ void GrpcRpcManager::MainWorkLoop() {
       MaybeCreatePoller();
 
       // Do actual work
-      DoWork();
+      DoWork(tag);
     }
 
     // Continue to loop if this thread can continue as a poller
@@ -158,7 +159,7 @@ void GrpcRpcManager::MainWorkLoop() {
     grpc::unique_lock<grpc::mutex> lock(mu_);
     num_threads_--;
     if (num_threads_ == 0) {
-      shutdown_cv_.notify_one();
+      shutdown_cv_.notify_all();
     }
   }
 
diff --git a/src/cpp/rpcmanager/grpc_rpc_manager.h b/src/cpp/rpcmanager/grpc_rpc_manager.h
index 5f89c1599d..a8cc6eb80f 100644
--- a/src/cpp/rpcmanager/grpc_rpc_manager.h
+++ b/src/cpp/rpcmanager/grpc_rpc_manager.h
@@ -50,12 +50,12 @@ class GrpcRpcManager {
   // This function MUST be called before using the object
   void Initialize();
 
-  virtual void PollForWork(bool& is_work_found) = 0;
-  virtual void DoWork() = 0;
+  virtual void PollForWork(bool& is_work_found, void **tag) = 0;
+  virtual void DoWork(void *tag) = 0;
 
-  // Use this for testing purposes only
+  // Use the following two functions for testing purposes only
   void Wait();
-  void Shutdown();
+  void ShutdownRpcManager();
 
  private:
   // Helper wrapper class around std::thread. This takes a GrpcRpcManager object