diff --git a/health/health.go b/health/health.go index 00007a306..de782dc51 100644 --- a/health/health.go +++ b/health/health.go @@ -3,6 +3,8 @@ package health import ( + "sync" + "golang.org/x/net/context" "google.golang.org/grpc" "google.golang.org/grpc/codes" @@ -11,13 +13,19 @@ import ( type HealthServer struct { // StatusMap stores the serving status of a service - StatusMap map[string]int32 + statusMap map[string]int32 +} + +func NewHealthServer() *HealthServer { + return &HealthServer{ + statusMap: make(map[string]int32), + } } func (s *HealthServer) Check(ctx context.Context, in *healthpb.HealthCheckRequest) (out *healthpb.HealthCheckResponse, err error) { service := in.Host + ":" + in.Service out = new(healthpb.HealthCheckResponse) - status, ok := s.StatusMap[service] + status, ok := s.statusMap[service] out.Status = healthpb.HealthCheckResponse_ServingStatus(status) if !ok { err = grpc.Errorf(codes.NotFound, "unknown service") @@ -31,5 +39,8 @@ func (s *HealthServer) Check(ctx context.Context, in *healthpb.HealthCheckReques // or insert a new service entry into the statusMap func (s *HealthServer) SetServingStatus(host string, service string, status int32) { service = host + ":" + service - s.StatusMap[service] = status + var mu sync.Mutex + mu.Lock() + s.statusMap[service] = status + mu.Unlock() } diff --git a/test/end2end_test.go b/test/end2end_test.go index 02950ff1d..f220ffcae 100644 --- a/test/end2end_test.go +++ b/test/end2end_test.go @@ -378,9 +378,7 @@ func TestHealthCheckOnSuccess(t *testing.T) { } func testHealthCheckOnSuccess(t *testing.T, e env) { - hs := &health.HealthServer{ - StatusMap: make(map[string]int32), - } + hs := health.NewHealthServer() hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", 1) s, cc := setUp(hs, math.MaxUint32, e) defer tearDown(s, cc) @@ -396,9 +394,7 @@ func TestHealthCheckOnFailure(t *testing.T) { } func testHealthCheckOnFailure(t *testing.T, e env) { - hs := &health.HealthServer{ - StatusMap: make(map[string]int32), - } + hs := health.NewHealthServer() hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", 1) s, cc := setUp(hs, math.MaxUint32, e) defer tearDown(s, cc) @@ -428,9 +424,7 @@ func TestHealthCheckNotFound(t *testing.T) { } func testHealthCheckNotFound(t *testing.T, e env) { - hs := &health.HealthServer{ - StatusMap: make(map[string]int32), - } + hs := health.NewHealthServer() s, cc := setUp(hs, math.MaxUint32, e) defer tearDown(s, cc) if _, err := healthCheck(1*time.Second, cc, "unregister_service"); err != grpc.Errorf(codes.NotFound, "unknown service") { @@ -445,9 +439,7 @@ func TestHealthCheckServing(t *testing.T) { } func testHealthCheckServing(t *testing.T, e env) { - hs := &health.HealthServer{ - StatusMap: make(map[string]int32), - } + hs := health.NewHealthServer() hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", 1) s, cc := setUp(hs, math.MaxUint32, e) defer tearDown(s, cc) @@ -467,9 +459,7 @@ func TestHealthCheckNotServing(t *testing.T) { } func testHealthCheckNotServing(t *testing.T, e env) { - hs := &health.HealthServer{ - StatusMap: make(map[string]int32), - } + hs := health.NewHealthServer() hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", 2) s, cc := setUp(hs, math.MaxUint32, e) defer tearDown(s, cc)