server.go

Documentation: google.golang.org/grpc/health

     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  // Package health provides a service that exposes server's health and it must be
    20  // imported to enable support for client-side health checks.
    21  package health
    22  
    23  import (
    24  	"context"
    25  	"sync"
    26  
    27  	"google.golang.org/grpc/codes"
    28  	healthgrpc "google.golang.org/grpc/health/grpc_health_v1"
    29  	healthpb "google.golang.org/grpc/health/grpc_health_v1"
    30  	"google.golang.org/grpc/status"
    31  )
    32  
    33  // Server implements `service Health`.
    34  type Server struct {
    35  	healthgrpc.UnimplementedHealthServer
    36  	mu sync.RWMutex
    37  	// If shutdown is true, it's expected all serving status is NOT_SERVING, and
    38  	// will stay in NOT_SERVING.
    39  	shutdown bool
    40  	// statusMap stores the serving status of the services this Server monitors.
    41  	statusMap map[string]healthpb.HealthCheckResponse_ServingStatus
    42  	updates   map[string]map[healthgrpc.Health_WatchServer]chan healthpb.HealthCheckResponse_ServingStatus
    43  }
    44  
    45  // NewServer returns a new Server.
    46  func NewServer() *Server {
    47  	return &Server{
    48  		statusMap: map[string]healthpb.HealthCheckResponse_ServingStatus{"": healthpb.HealthCheckResponse_SERVING},
    49  		updates:   make(map[string]map[healthgrpc.Health_WatchServer]chan healthpb.HealthCheckResponse_ServingStatus),
    50  	}
    51  }
    52  
    53  // Check implements `service Health`.
    54  func (s *Server) Check(ctx context.Context, in *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
    55  	s.mu.RLock()
    56  	defer s.mu.RUnlock()
    57  	if servingStatus, ok := s.statusMap[in.Service]; ok {
    58  		return &healthpb.HealthCheckResponse{
    59  			Status: servingStatus,
    60  		}, nil
    61  	}
    62  	return nil, status.Error(codes.NotFound, "unknown service")
    63  }
    64  
    65  // Watch implements `service Health`.
    66  func (s *Server) Watch(in *healthpb.HealthCheckRequest, stream healthgrpc.Health_WatchServer) error {
    67  	service := in.Service
    68  	// update channel is used for getting service status updates.
    69  	update := make(chan healthpb.HealthCheckResponse_ServingStatus, 1)
    70  	s.mu.Lock()
    71  	// Puts the initial status to the channel.
    72  	if servingStatus, ok := s.statusMap[service]; ok {
    73  		update <- servingStatus
    74  	} else {
    75  		update <- healthpb.HealthCheckResponse_SERVICE_UNKNOWN
    76  	}
    77  
    78  	// Registers the update channel to the correct place in the updates map.
    79  	if _, ok := s.updates[service]; !ok {
    80  		s.updates[service] = make(map[healthgrpc.Health_WatchServer]chan healthpb.HealthCheckResponse_ServingStatus)
    81  	}
    82  	s.updates[service][stream] = update
    83  	defer func() {
    84  		s.mu.Lock()
    85  		delete(s.updates[service], stream)
    86  		s.mu.Unlock()
    87  	}()
    88  	s.mu.Unlock()
    89  
    90  	var lastSentStatus healthpb.HealthCheckResponse_ServingStatus = -1
    91  	for {
    92  		select {
    93  		// Status updated. Sends the up-to-date status to the client.
    94  		case servingStatus := <-update:
    95  			if lastSentStatus == servingStatus {
    96  				continue
    97  			}
    98  			lastSentStatus = servingStatus
    99  			err := stream.Send(&healthpb.HealthCheckResponse{Status: servingStatus})
   100  			if err != nil {
   101  				return status.Error(codes.Canceled, "Stream has ended.")
   102  			}
   103  		// Context done. Removes the update channel from the updates map.
   104  		case <-stream.Context().Done():
   105  			return status.Error(codes.Canceled, "Stream has ended.")
   106  		}
   107  	}
   108  }
   109  
   110  // SetServingStatus is called when need to reset the serving status of a service
   111  // or insert a new service entry into the statusMap.
   112  func (s *Server) SetServingStatus(service string, servingStatus healthpb.HealthCheckResponse_ServingStatus) {
   113  	s.mu.Lock()
   114  	defer s.mu.Unlock()
   115  	if s.shutdown {
   116  		logger.Infof("health: status changing for %s to %v is ignored because health service is shutdown", service, servingStatus)
   117  		return
   118  	}
   119  
   120  	s.setServingStatusLocked(service, servingStatus)
   121  }
   122  
   123  func (s *Server) setServingStatusLocked(service string, servingStatus healthpb.HealthCheckResponse_ServingStatus) {
   124  	s.statusMap[service] = servingStatus
   125  	for _, update := range s.updates[service] {
   126  		// Clears previous updates, that are not sent to the client, from the channel.
   127  		// This can happen if the client is not reading and the server gets flow control limited.
   128  		select {
   129  		case <-update:
   130  		default:
   131  		}
   132  		// Puts the most recent update to the channel.
   133  		update <- servingStatus
   134  	}
   135  }
   136  
   137  // Shutdown sets all serving status to NOT_SERVING, and configures the server to
   138  // ignore all future status changes.
   139  //
   140  // This changes serving status for all services. To set status for a particular
   141  // services, call SetServingStatus().
   142  func (s *Server) Shutdown() {
   143  	s.mu.Lock()
   144  	defer s.mu.Unlock()
   145  	s.shutdown = true
   146  	for service := range s.statusMap {
   147  		s.setServingStatusLocked(service, healthpb.HealthCheckResponse_NOT_SERVING)
   148  	}
   149  }
   150  
   151  // Resume sets all serving status to SERVING, and configures the server to
   152  // accept all future status changes.
   153  //
   154  // This changes serving status for all services. To set status for a particular
   155  // services, call SetServingStatus().
   156  func (s *Server) Resume() {
   157  	s.mu.Lock()
   158  	defer s.mu.Unlock()
   159  	s.shutdown = false
   160  	for service := range s.statusMap {
   161  		s.setServingStatusLocked(service, healthpb.HealthCheckResponse_SERVING)
   162  	}
   163  }
   164
View as plain text