diff --git a/health/grpc_health_v1alpha/health.proto b/health/grpc_health_v1alpha/health.proto index f2c7efcf799b..de7c04580efb 100644 --- a/health/grpc_health_v1alpha/health.proto +++ b/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 { - enum ServingStatus { - UNKNOWN = 0; - SERVING = 1; - NOT_SERVING = 2; - } - ServingStatus status = 1; - } - +message HealthCheckResponse { + enum ServingStatus { + UNKNOWN = 0; + SERVING = 1; + NOT_SERVING = 2; + } + ServingStatus status = 1; + } service HealthCheck{ - rpc Check( HealthCheckRequest) returns ( HealthCheckResponse); -} + rpc Check( HealthCheckRequest) returns ( HealthCheckResponse); + } diff --git a/health/health.go b/health/health.go index 66081d60e28b..00007a306ebf 100644 --- a/health/health.go +++ b/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) { - service = ":" + service +// SetServingStatus is called when need to reset the serving status of a 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 } diff --git a/test/end2end_test.go b/test/end2end_test.go index 65dee8994300..b4f67957faff 100644 --- a/test/end2end_test.go +++ b/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 { - t.Fatalf("HealthCheck(_)=_, %v, want ", err) + if _, err := healthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck"); err != nil { + t.Fatalf("TestHealthCheck(_)=_, %v, want ", 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") { - t.Fatalf("HealthCheck(_)=_, %v, want error code %d", err, codes.DeadlineExceeded) + if _, err := healthCheck(0*time.Second, cc, "grpc.health.v1alpha.HealthCheck"); err != grpc.Errorf(codes.DeadlineExceeded, "context deadline exceeded") { + 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") { - t.Fatalf("HealthCheck(_)=_, %v, want error code %d", err, codes.Unimplemented) + if _, err := healthCheck(1*time.Second, cc, ""); err != grpc.Errorf(codes.Unimplemented, "unknown service grpc.health.v1alpha.HealthCheck") { + 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") { - t.Fatalf("HealthCheck(_)=_, %v, want error code %d", err, codes.NotFound) + if _, err := healthCheck(1*time.Second, cc, "unregister_service"); err != grpc.Errorf(codes.NotFound, "unknown service") { + 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 _,", err) + t.Fatalf("TestHealthCheck(_)=_, %v, want _,", 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 _,", err) + t.Fatalf("TestHealthCheck(_)=_, %v, want _,", err) } if out.Status != healthpb.HealthCheckResponse_NOT_SERVING { t.Fatalf("Got the serving status %v, want NOT_SERVING ")