vendor/google.golang.org/grpc/health/health.go

   1 /*
   2  *
   3  * Copyright 2017 gRPC authors.
   4  *
   5  * Licensed under the Apache License, Version 2.0 (the "License");
   6  * you may not use this file except in compliance with the License.
   7  * You may obtain a copy of the License at
   8  *
   9  *     http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  *
  17  */
  18
  19 //go:generate protoc --go_out=plugins=grpc:. grpc_health_v1/health.proto
  20
  21 // Package health provides some utility functions to health-check a server. The implementation
  22 // is based on protobuf. Users need to write their own implementations if other IDLs are used.
  23 package health
  24
  25 import (
  26         "sync"
  27
  28         "golang.org/x/net/context"
  29         "google.golang.org/grpc"
  30         "google.golang.org/grpc/codes"
  31         healthpb "google.golang.org/grpc/health/grpc_health_v1"
  32 )
  33
  34 // Server implements `service Health`.
  35 type Server struct {
  36         mu sync.Mutex
  37         // statusMap stores the serving status of the services this Server monitors.
  38         statusMap map[string]healthpb.HealthCheckResponse_ServingStatus
  39 }
  40
  41 // NewServer returns a new Server.
  42 func NewServer() *Server {
  43         return &Server{
  44                 statusMap: make(map[string]healthpb.HealthCheckResponse_ServingStatus),
  45         }
  46 }
  47
  48 // Check implements `service Health`.
  49 func (s *Server) Check(ctx context.Context, in *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
  50         s.mu.Lock()
  51         defer s.mu.Unlock()
  52         if in.Service == "" {
  53                 // check the server overall health status.
  54                 return &healthpb.HealthCheckResponse{
  55                         Status: healthpb.HealthCheckResponse_SERVING,
  56                 }, nil
  57         }
  58         if status, ok := s.statusMap[in.Service]; ok {
  59                 return &healthpb.HealthCheckResponse{
  60                         Status: status,
  61                 }, nil
  62         }
  63         return nil, grpc.Errorf(codes.NotFound, "unknown service")
  64 }
  65
  66 // SetServingStatus is called when need to reset the serving status of a service
  67 // or insert a new service entry into the statusMap.
  68 func (s *Server) SetServingStatus(service string, status healthpb.HealthCheckResponse_ServingStatus) {
  69         s.mu.Lock()
  70         s.statusMap[service] = status
  71         s.mu.Unlock()
  72 }