1 files changed, 19 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..44aebd2f9d 100644
--- a/src/cpp/server/health/default_health_check_service.cc
+++ b/src/cpp/server/health/default_health_check_service.cc
@@ -42,18 +42,37 @@ DefaultHealthCheckService::DefaultHealthCheckService() {
 void DefaultHealthCheckService::SetServingStatus(
     const grpc::string& service_name, bool serving) {
   std::unique_lock<std::mutex> lock(mu_);
+  if (shutdown_) {
+    // Set to NOT_SERVING in case service_name is not in the map.
+    serving = false;
+  }
   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 {