health: resume health server

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
 	}
 
@@ -143,8 +141,19 @@ func (s *Server) setServingStatusLocked(service string, servingStatus healthpb.H
 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)
 	}
 }
+
+// Restart sets all serving status to SERVING, and configures the server to
+// accept all future status changes.
+func (s *Server) Restart() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.shutdown = false
+	for service := range s.statusMap {
+		s.setServingStatusLocked(service, healthpb.HealthCheckResponse_SERVING)
+	}
+}

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.Restart()
+	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)
+	}
 }