minor changes

This commit is contained in:
yangzhouhan
2015-07-22 16:04:53 -07:00
parent 2fdea48f1f
commit 34f97a03ae
3 changed files with 44 additions and 48 deletions

View File

@ -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);
}

View File

@ -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
}

View File

@ -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 <nil>", err)
if _, err := healthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck"); err != nil {
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") {
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 _,<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 ")