Skip to content

Commit

Permalink
minor changes
Browse files Browse the repository at this point in the history
  • Loading branch information
@yangzhouhan
yangzhouhan committed Jul 22, 2015
1 parent 2fdea48 commit 34f97a0
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 48 deletions.
23 changes: 11 additions & 12 deletions health/grpc_health_v1alpha/health.proto
View file
Original file line number Diff line number Diff line change
@@ -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);
}
22 changes: 6 additions & 16 deletions health/health.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
}
47 changes: 27 additions & 20 deletions test/end2end_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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 <nil>", err)
if _, err := healthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck"); err != nil {
t.Fatalf("TestHealthCheck(_)=_, %v, want <nil>", err)
}
}

Expand All @@ -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)
}
}

Expand All @@ -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)
}
}

Expand All @@ -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)
}
}

Expand All @@ -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)
Expand All @@ -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 ")
Expand Down

0 comments on commit 34f97a0

Please sign in to comment.