diff --git a/health/server.go b/health/server.go index 3e3f2e3f26c2..c79f9d2ab80c 100644 --- a/health/server.go +++ b/health/server.go @@ -27,19 +27,18 @@ import ( "sync" "google.golang.org/grpc/codes" + "google.golang.org/grpc/grpclog" healthgrpc "google.golang.org/grpc/health/grpc_health_v1" healthpb "google.golang.org/grpc/health/grpc_health_v1" - "google.golang.org/grpc/internal/grpcsync" "google.golang.org/grpc/status" ) // Server implements `service Health`. type Server struct { - // If shutdownEvent has fired, it's expected all serving status is - // NOT_SERVING, and will stay in NOT_SERVING. - shutdownEvent *grpcsync.Event - mu sync.Mutex + // If shutdown is true, it's expected all serving status is NOT_SERVING, and + // will stay in NOT_SERVING. + shutdown bool // statusMap stores the serving status of the services this Server monitors. statusMap map[string]healthpb.HealthCheckResponse_ServingStatus updates map[string]map[healthgrpc.Health_WatchServer]chan healthpb.HealthCheckResponse_ServingStatus @@ -48,8 +47,6 @@ type Server struct { // NewServer returns a new Server. func NewServer() *Server { return &Server{ - shutdownEvent: grpcsync.NewEvent(), - statusMap: map[string]healthpb.HealthCheckResponse_ServingStatus{"": healthpb.HealthCheckResponse_SERVING}, updates: make(map[string]map[healthgrpc.Health_WatchServer]chan healthpb.HealthCheckResponse_ServingStatus), } @@ -117,7 +114,8 @@ func (s *Server) Watch(in *healthpb.HealthCheckRequest, stream healthgrpc.Health func (s *Server) SetServingStatus(service string, servingStatus healthpb.HealthCheckResponse_ServingStatus) { s.mu.Lock() defer s.mu.Unlock() - if s.shutdownEvent.HasFired() { + if s.shutdown { + grpclog.Infof("health: status changing for %s to %v is ignored because health service is shutdown", service, servingStatus) return } @@ -140,11 +138,28 @@ func (s *Server) setServingStatusLocked(service string, servingStatus healthpb.H // Shutdown sets all serving status to NOT_SERVING, and configures the server to // ignore all future status changes. +// +// This changes serving status for all services. To set status for a perticular +// services, call SetServingStatus(). func (s *Server) Shutdown() { s.mu.Lock() defer s.mu.Unlock() - s.shutdownEvent.Fire() + s.shutdown = true for service := range s.statusMap { s.setServingStatusLocked(service, healthpb.HealthCheckResponse_NOT_SERVING) } } + +// Resume sets all serving status to SERVING, and configures the server to +// accept all future status changes. +// +// This changes serving status for all services. To set status for a perticular +// services, call SetServingStatus(). +func (s *Server) Resume() { + s.mu.Lock() + defer s.mu.Unlock() + s.shutdown = false + for service := range s.statusMap { + s.setServingStatusLocked(service, healthpb.HealthCheckResponse_SERVING) + } +} diff --git a/health/server_internal_test.go b/health/server_internal_test.go index fc137ba63474..735faa7892d2 100644 --- a/health/server_internal_test.go +++ b/health/server_internal_test.go @@ -59,4 +59,16 @@ func TestShutdown(t *testing.T) { if status != healthpb.HealthCheckResponse_NOT_SERVING { t.Fatalf("status for %s is %v, want %v", testService, status, healthpb.HealthCheckResponse_NOT_SERVING) } + + s.Resume() + status = s.statusMap[testService] + if status != healthpb.HealthCheckResponse_SERVING { + t.Fatalf("status for %s is %v, want %v", testService, status, healthpb.HealthCheckResponse_SERVING) + } + + s.SetServingStatus(testService, healthpb.HealthCheckResponse_NOT_SERVING) + status = s.statusMap[testService] + if status != healthpb.HealthCheckResponse_NOT_SERVING { + t.Fatalf("status for %s is %v, want %v", testService, status, healthpb.HealthCheckResponse_NOT_SERVING) + } }