add mutex and create newHealthServer Function

This commit is contained in:
yangzhouhan 2015-07-22 16:26:16 -07:00
parent 901c3e574c
commit dc36cdbfdd
2 changed files with 19 additions and 18 deletions

View File

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

View File

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