Add a Shutdown call to HealthCheckServiceInterface

2 files changed, 21 insertions, 0 deletions

diff --git a/src/cpp/server/health/default_health_check_service.cc b/src/cpp/server/health/default_health_check_service.cc
index c951c69d51..db6286d240 100644
--- a/src/cpp/server/health/default_health_check_service.cc
+++ b/src/cpp/server/health/default_health_check_service.cc
@@ -42,18 +42,36 @@ DefaultHealthCheckService::DefaultHealthCheckService() {
 void DefaultHealthCheckService::SetServingStatus(
     const grpc::string& service_name, bool serving) {
   std::unique_lock<std::mutex> lock(mu_);
+  if (shutdown_) {
+    return;
+  }
   services_map_[service_name].SetServingStatus(serving ? SERVING : NOT_SERVING);
 }
 
 void DefaultHealthCheckService::SetServingStatus(bool serving) {
   const ServingStatus status = serving ? SERVING : NOT_SERVING;
   std::unique_lock<std::mutex> lock(mu_);
+  if (shutdown_) {
+    return;
+  }
   for (auto& p : services_map_) {
     ServiceData& service_data = p.second;
     service_data.SetServingStatus(status);
   }
 }
 
+void DefaultHealthCheckService::Shutdown() {
+  std::unique_lock<std::mutex> lock(mu_);
+  if (shutdown_) {
+    return;
+  }
+  shutdown_ = true;
+  for (auto& p : services_map_) {
+    ServiceData& service_data = p.second;
+    service_data.SetServingStatus(NOT_SERVING);
+  }
+}
+
 DefaultHealthCheckService::ServingStatus
 DefaultHealthCheckService::GetServingStatus(
     const grpc::string& service_name) const {
diff --git a/src/cpp/server/health/default_health_check_service.h b/src/cpp/server/health/default_health_check_service.h
index 450bd543f5..9551cd2e2c 100644
--- a/src/cpp/server/health/default_health_check_service.h
+++ b/src/cpp/server/health/default_health_check_service.h
@@ -237,6 +237,8 @@ class DefaultHealthCheckService final : public HealthCheckServiceInterface {
                         bool serving) override;
   void SetServingStatus(bool serving) override;
 
+  void Shutdown() override;
+
   ServingStatus GetServingStatus(const grpc::string& service_name) const;
 
   HealthCheckServiceImpl* GetHealthCheckService(
@@ -272,6 +274,7 @@ class DefaultHealthCheckService final : public HealthCheckServiceInterface {
       const std::shared_ptr<HealthCheckServiceImpl::CallHandler>& handler);
 
   mutable std::mutex mu_;
+  bool shutdown_ = false;                             // Guarded by mu_.
   std::map<grpc::string, ServiceData> services_map_;  // Guarded by mu_.
   std::unique_ptr<HealthCheckServiceImpl> impl_;
 };