mirror of https://github.com/grpc/grpc-go.git
add mutex and create newHealthServer Function
This commit is contained in:
parent
901c3e574c
commit
dc36cdbfdd
|
@ -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()
|
||||
}
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue