From dc36cdbfdd56bcd133e87cb0c901d7720fb6486a Mon Sep 17 00:00:00 2001
From: yangzhouhan <yangzhouhan@gmail.com>
Date: Wed, 22 Jul 2015 16:26:16 -0700
Subject: [PATCH] add mutex and create newHealthServer Function

---
 health/health.go     | 17 ++++++++++++++---
 test/end2end_test.go | 20 +++++---------------
 2 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/health/health.go b/health/health.go
index 00007a30..de782dc5 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 02950ff1..f220ffca 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)