opensource/grpc-go
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add mutex and create newHealthServer Function

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

17
health/health.go
View File

@ -3,6 +3,8 @@
package health package health
import ( import (
"sync"
"golang.org/x/net/context" "golang.org/x/net/context"
"google.golang.org/grpc" "google.golang.org/grpc"
"google.golang.org/grpc/codes" "google.golang.org/grpc/codes"
@ -11,13 +13,19 @@ import (
type HealthServer struct { type HealthServer struct {
// StatusMap stores the serving status of a service // 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) { func (s *HealthServer) Check(ctx context.Context, in *healthpb.HealthCheckRequest) (out *healthpb.HealthCheckResponse, err error) {
service := in.Host + ":" + in.Service service := in.Host + ":" + in.Service
out = new(healthpb.HealthCheckResponse) out = new(healthpb.HealthCheckResponse)
status, ok := s.StatusMap[service] status, ok := s.statusMap[service]
out.Status = healthpb.HealthCheckResponse_ServingStatus(status) out.Status = healthpb.HealthCheckResponse_ServingStatus(status)
if !ok { if !ok {
err = grpc.Errorf(codes.NotFound, "unknown service") 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 // or insert a new service entry into the statusMap
func (s *HealthServer) SetServingStatus(host string, service string, status int32) { func (s *HealthServer) SetServingStatus(host string, service string, status int32) {
service = host + ":" + service service = host + ":" + service
s.StatusMap[service] = status var mu sync.Mutex
mu.Lock()
s.statusMap[service] = status
mu.Unlock()
} }

20
test/end2end_test.go
View File

@ -378,9 +378,7 @@ func TestHealthCheckOnSuccess(t *testing.T) {
} }
func testHealthCheckOnSuccess(t *testing.T, e env) { func testHealthCheckOnSuccess(t *testing.T, e env) {
hs := &health.HealthServer{ hs := health.NewHealthServer()
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) s, cc := setUp(hs, math.MaxUint32, e)
defer tearDown(s, cc) defer tearDown(s, cc)
@ -396,9 +394,7 @@ func TestHealthCheckOnFailure(t *testing.T) {
} }
func testHealthCheckOnFailure(t *testing.T, e env) { func testHealthCheckOnFailure(t *testing.T, e env) {
hs := &health.HealthServer{ hs := health.NewHealthServer()
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) s, cc := setUp(hs, math.MaxUint32, e)
defer tearDown(s, cc) defer tearDown(s, cc)
@ -428,9 +424,7 @@ func TestHealthCheckNotFound(t *testing.T) {
} }
func testHealthCheckNotFound(t *testing.T, e env) { func testHealthCheckNotFound(t *testing.T, e env) {
hs := &health.HealthServer{ hs := health.NewHealthServer()
StatusMap: make(map[string]int32),
}
s, cc := setUp(hs, math.MaxUint32, e) s, cc := setUp(hs, math.MaxUint32, e)
defer tearDown(s, cc) defer tearDown(s, cc)
if _, err := 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") {
@ -445,9 +439,7 @@ func TestHealthCheckServing(t *testing.T) {
} }
func testHealthCheckServing(t *testing.T, e env) { func testHealthCheckServing(t *testing.T, e env) {
hs := &health.HealthServer{ hs := health.NewHealthServer()
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) s, cc := setUp(hs, math.MaxUint32, e)
defer tearDown(s, cc) defer tearDown(s, cc)
@ -467,9 +459,7 @@ func TestHealthCheckNotServing(t *testing.T) {
} }
func testHealthCheckNotServing(t *testing.T, e env) { func testHealthCheckNotServing(t *testing.T, e env) {
hs := &health.HealthServer{ hs := health.NewHealthServer()
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) s, cc := setUp(hs, math.MaxUint32, e)
defer tearDown(s, cc) defer tearDown(s, cc)