github.com/stackdocker/rkt@v0.10.1-0.20151109095037-1aa827478248/Godeps/_workspace/src/google.golang.org/grpc/health/health.go (about)

     1  // Package health provides some utility functions to health-check a server. The implementation
     2  // is based on protobuf. Users need to write their own implementations if other IDLs are used.
     3  package health
     4  
     5  import (
     6  	"sync"
     7  
     8  	"github.com/coreos/rkt/Godeps/_workspace/src/golang.org/x/net/context"
     9  	"github.com/coreos/rkt/Godeps/_workspace/src/google.golang.org/grpc"
    10  	"github.com/coreos/rkt/Godeps/_workspace/src/google.golang.org/grpc/codes"
    11  	healthpb "github.com/coreos/rkt/Godeps/_workspace/src/google.golang.org/grpc/health/grpc_health_v1alpha"
    12  )
    13  
    14  type HealthServer struct {
    15  	mu sync.Mutex
    16  	// statusMap stores the serving status of the services this HealthServer monitors.
    17  	statusMap map[string]healthpb.HealthCheckResponse_ServingStatus
    18  }
    19  
    20  func NewHealthServer() *HealthServer {
    21  	return &HealthServer{
    22  		statusMap: make(map[string]healthpb.HealthCheckResponse_ServingStatus),
    23  	}
    24  }
    25  
    26  func (s *HealthServer) Check(ctx context.Context, in *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
    27  	s.mu.Lock()
    28  	defer s.mu.Unlock()
    29  	if in.Service == "" {
    30  		// check the server overall health status.
    31  		return &healthpb.HealthCheckResponse{
    32  			Status: healthpb.HealthCheckResponse_SERVING,
    33  		}, nil
    34  	}
    35  	if status, ok := s.statusMap[in.Service]; ok {
    36  		return &healthpb.HealthCheckResponse{
    37  			Status: status,
    38  		}, nil
    39  	}
    40  	return nil, grpc.Errorf(codes.NotFound, "unknown service")
    41  }
    42  
    43  // SetServingStatus is called when need to reset the serving status of a service
    44  // or insert a new service entry into the statusMap.
    45  func (s *HealthServer) SetServingStatus(service string, status healthpb.HealthCheckResponse_ServingStatus) {
    46  	s.mu.Lock()
    47  	s.statusMap[service] = status
    48  	s.mu.Unlock()
    49  }