minor changes

health/grpc_health_v1alpha/health.proto

@@ -1,20 +1,19 @@
 syntax = "proto3";
+
 package grpc.health.v1alpha;
 
 message HealthCheckRequest {
   string host = 1;
   string service = 2;
   }
-
+message HealthCheckResponse {
-	message HealthCheckResponse {
+  enum ServingStatus {
-	  enum ServingStatus {
+    UNKNOWN = 0;
-	          UNKNOWN = 0;
+	SERVING = 1;
-		      SERVING = 1;
+	NOT_SERVING = 2;
-			  NOT_SERVING = 2;
+	}
-			  }
+  ServingStatus status = 1;
-			  ServingStatus status = 1;
+  }
-		  }
-
 service HealthCheck{
-	rpc Check( HealthCheckRequest) returns ( HealthCheckResponse);
+  rpc Check( HealthCheckRequest) returns ( HealthCheckResponse);
-}
+  } 

health/health.go

@@ -3,31 +3,19 @@
 package health
 
 import (
-	"time"
-
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	healthpb "google.golang.org/grpc/health/grpc_health_v1alpha"
 )
 
-// HealthCheck is the client side function to health-check a server
-func HealthCheck(t time.Duration, cc *grpc.ClientConn, service_name string) (*healthpb.HealthCheckResponse, error) {
-	ctx, _ := context.WithTimeout(context.Background(), t)
-	hc := healthpb.NewHealthCheckClient(cc)
-	req := new(healthpb.HealthCheckRequest)
-	req.Host = ""
-	req.Service = service_name
-	out, err := hc.Check(ctx, req)
-	return out, err
-}
-
 type HealthServer struct {
+	// StatusMap stores the serving status of a service
 	StatusMap map[string]int32
 }
 
 func (s *HealthServer) Check(ctx context.Context, in *healthpb.HealthCheckRequest) (out *healthpb.HealthCheckResponse, err error) {
-	service := ":" + in.Service
+	service := in.Host + ":" + in.Service
 	out = new(healthpb.HealthCheckResponse)
 	status, ok := s.StatusMap[service]
 	out.Status = healthpb.HealthCheckResponse_ServingStatus(status)
@@ -39,7 +27,9 @@ func (s *HealthServer) Check(ctx context.Context, in *healthpb.HealthCheckReques
 	return out, err
 }
 
-func (s *HealthServer) SetServingStatus(service string, status int32) {
+// SetServingStatus is called when need to reset the serving status of a service
-	service = ":" + service
+// 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
 }

test/end2end_test.go

@@ -93,6 +93,17 @@ func newPayload(t testpb.PayloadType, size int32) *testpb.Payload {
 	}
 }
 
+func healthCheck(t time.Duration, cc *grpc.ClientConn, serviceName string) (*healthpb.HealthCheckResponse, error) {
+	ctx, _ := context.WithTimeout(context.Background(), t)
+	hc := healthpb.NewHealthCheckClient(cc)
+	req := &healthpb.HealthCheckRequest{
+		Host:    "",
+		Service: serviceName,
+	}
+	out, err := hc.Check(ctx, req)
+	return out, err
+}
+
 func (s *testServer) UnaryCall(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
 	md, ok := metadata.FromContext(ctx)
 	if ok {
@@ -371,11 +382,11 @@ func testHealthCheckOnSuccess(t *testing.T, e env) {
 	hs := &health.HealthServer{
 		StatusMap: make(map[string]int32),
 	}
-	hs.SetServingStatus("grpc.health.v1alpha.HealthCheck", 1)
+	hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", 1)
 	s, cc := setUp(hs, math.MaxUint32, e)
 	defer tearDown(s, cc)
-	if _, err := health.HealthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck"); err != nil {
+	if _, err := healthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck"); err != nil {
-		t.Fatalf("HealthCheck(_)=_, %v, want <nil>", err)
+		t.Fatalf("TestHealthCheck(_)=_, %v, want <nil>", err)
 	}
 }
 
@@ -390,11 +401,11 @@ func testHealthCheckOnFailure(t *testing.T, e env) {
 	hs := &health.HealthServer{
 		StatusMap: make(map[string]int32),
 	}
-	hs.SetServingStatus("grpc.health.v1alpha.HealthCheck", 1)
+	hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", 1)
 	s, cc := setUp(hs, math.MaxUint32, e)
 	defer tearDown(s, cc)
-	if _, err := health.HealthCheck(0*time.Second, cc, "grpc.health.v1alpha.HealthCheck"); err != grpc.Errorf(codes.DeadlineExceeded, "context deadline exceeded") {
+	if _, err := healthCheck(0*time.Second, cc, "grpc.health.v1alpha.HealthCheck"); err != grpc.Errorf(codes.DeadlineExceeded, "context deadline exceeded") {
-		t.Fatalf("HealthCheck(_)=_, %v, want error code %d", err, codes.DeadlineExceeded)
+		t.Fatalf("TestHealthCheck(_)=_, %v, want error code %d", err, codes.DeadlineExceeded)
 	}
 }
 
@@ -407,8 +418,8 @@ func TestHealthCheckOff(t *testing.T) {
 func testHealthCheckOff(t *testing.T, e env) {
 	s, cc := setUp(nil, math.MaxUint32, e)
 	defer tearDown(s, cc)
-	if _, err := health.HealthCheck(1*time.Second, cc, ""); err != grpc.Errorf(codes.Unimplemented, "unknown service grpc.health.v1alpha.HealthCheck") {
+	if _, err := healthCheck(1*time.Second, cc, ""); err != grpc.Errorf(codes.Unimplemented, "unknown service grpc.health.v1alpha.HealthCheck") {
-		t.Fatalf("HealthCheck(_)=_, %v, want error code %d", err, codes.Unimplemented)
+		t.Fatalf("TestHealthCheck(_)=_, %v, want error code %d", err, codes.Unimplemented)
 	}
 }
 
@@ -419,15 +430,13 @@ func TestHealthCheckNotFound(t *testing.T) {
 }
 
 func testHealthCheckNotFound(t *testing.T, e env) {
-
 	hs := &health.HealthServer{
 		StatusMap: make(map[string]int32),
 	}
-
 	s, cc := setUp(hs, math.MaxUint32, e)
 	defer tearDown(s, cc)
-	if _, err := health.HealthCheck(1*time.Second, cc, "unregister_service"); err != grpc.Errorf(codes.NotFound, "unknown service") {
+	if _, err := healthCheck(1*time.Second, cc, "unregister_service"); err != grpc.Errorf(codes.NotFound, "unknown service") {
-		t.Fatalf("HealthCheck(_)=_, %v, want error code %d", err, codes.NotFound)
+		t.Fatalf("TestHealthCheck(_)=_, %v, want error code %d", err, codes.NotFound)
 	}
 }
 
@@ -438,16 +447,15 @@ func TestHealthCheckServing(t *testing.T) {
 }
 
 func testHealthCheckServing(t *testing.T, e env) {
-
 	hs := &health.HealthServer{
 		StatusMap: make(map[string]int32),
 	}
-	hs.SetServingStatus("grpc.health.v1alpha.HealthCheck", 1)
+	hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", 1)
 	s, cc := setUp(hs, math.MaxUint32, e)
 	defer tearDown(s, cc)
-	out, err := health.HealthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck")
+	out, err := healthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck")
 	if err != nil {
-		t.Fatalf("HealthCheck(_)=_, %v, want _,<nil>", err)
+		t.Fatalf("TestHealthCheck(_)=_, %v, want _,<nil>", err)
 	}
 	if out.Status != healthpb.HealthCheckResponse_SERVING {
 		t.Fatalf("Got the serving status %v, want SERVING", out.Status)
@@ -461,16 +469,15 @@ func TestHealthCheckNotServing(t *testing.T) {
 }
 
 func testHealthCheckNotServing(t *testing.T, e env) {
-
 	hs := &health.HealthServer{
 		StatusMap: make(map[string]int32),
 	}
-	hs.SetServingStatus("grpc.health.v1alpha.HealthCheck", 2)
+	hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", 2)
 	s, cc := setUp(hs, math.MaxUint32, e)
 	defer tearDown(s, cc)
-	out, err := health.HealthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck")
+	out, err := healthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck")
 	if err != nil {
-		t.Fatalf("HealthCheck(_)=_, %v, want _,<nil>", err)
+		t.Fatalf("TestHealthCheck(_)=_, %v, want _,<nil>", err)
 	}
 	if out.Status != healthpb.HealthCheckResponse_NOT_SERVING {
 		t.Fatalf("Got the serving status %v, want NOT_SERVING ")