53 lines
1.6 KiB
Go
53 lines
1.6 KiB
Go
|
// Package health provides some utility functions to health-check a server. The implementation
|
||
|
|
// is based on protobuf. Users need to write their own implementations if other IDLs are used.
|
||
|
|
package health
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"golang.org/x/net/context"
|
||
|
|
"google.golang.org/grpc"
|
||
|
|
"google.golang.org/grpc/codes"
|
||
|
|
healthpb "google.golang.org/grpc/health/grpc_health_v1"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Server implements `service Health`.
|
||
|
|
type Server struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
// statusMap stores the serving status of the services this Server monitors.
|
||
|
|
statusMap map[string]healthpb.HealthCheckResponse_ServingStatus
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewServer returns a new Server.
|
||
|
|
func NewServer() *Server {
|
||
|
|
return &Server{
|
||
|
|
statusMap: make(map[string]healthpb.HealthCheckResponse_ServingStatus),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Check implements `service Health`.
|
||
|
|
func (s *Server) Check(ctx context.Context, in *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
|
||
|
|
s.mu.Lock()
|
||
|
|
defer s.mu.Unlock()
|
||
|
|
if in.Service == "" {
|
||
|
|
// check the server overall health status.
|
||
|
|
return &healthpb.HealthCheckResponse{
|
||
|
|
Status: healthpb.HealthCheckResponse_SERVING,
|
||
|
|
}, nil
|
||
|
|
}
|
||
|
|
if status, ok := s.statusMap[in.Service]; ok {
|
||
|
|
return &healthpb.HealthCheckResponse{
|
||
|
|
Status: status,
|
||
|
|
}, nil
|
||
|
|
}
|
||
|
|
return nil, grpc.Errorf(codes.NotFound, "unknown 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 *Server) SetServingStatus(service string, status healthpb.HealthCheckResponse_ServingStatus) {
|
||
|
|
s.mu.Lock()
|
||
|
|
s.statusMap[service] = status
|
||
|
|
s.mu.Unlock()
|
||
|
|
}
|