Repository: grpc/grpc-go
File: health/server.go (Go)

Code before:
func NewServer() *Server {
	return &Server{
		statusMap: map[string]healthpb.HealthCheckResponse_ServingStatus{"": healthpb.HealthCheckResponse_SERVING},
		updates:   make(map[string]map[healthgrpc.Health_WatchServer]chan healthpb.HealthCheckResponse_ServingStatus),
	}
}

// Check implements `service Health`.
func (s *Server) Check(_ context.Context, in *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if servingStatus, ok := s.statusMap[in.Service]; ok {
		return &healthpb.HealthCheckResponse{
			Status: servingStatus,
		}, nil
	}
	return nil, status.Error(codes.NotFound, "unknown service")
}

// List implements `service Health`.
func (s *Server) List(_ context.Context, _ *healthpb.HealthListRequest) (*healthpb.HealthListResponse, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if len(s.statusMap) > maxAllowedServices {
		return nil, status.Error(codes.ResourceExhausted, "server health list exceeds maximum capacity (100)")
	}

	statusMap := make(map[string]*healthpb.HealthCheckResponse, len(s.statusMap))
	for k, v := range s.statusMap {
		statusMap[k] = &healthpb.HealthCheckResponse{Status: v}
	}

	return &healthpb.HealthListResponse{Statuses: statusMap}, nil
}

// Watch implements `service Health`.
func (s *Server) Watch(in *healthpb.HealthCheckRequest, stream healthgrpc.Health_WatchServer) error {
	service := in.Service
	// update channel is used for getting service status updates.
	update := make(chan healthpb.HealthCheckResponse_ServingStatus, 1)
	s.mu.Lock()
	// Puts the initial status to the channel.
	if servingStatus, ok := s.statusMap[service]; ok {
		update <- servingStatus
	} else {
		update <- healthpb.HealthCheckResponse_SERVICE_UNKNOWN
	}

	// Registers the update channel to the correct place in the updates map.
	if _, ok := s.updates[service]; !ok {

Review comment (suggestion):
```suggestion
		return nil, status.Error(codes.ResourceExhausted, "server health list exceeds maximum capacity: %d", maxAllowedServices)
```

Code after:
// NewServer returns a new Server.
func NewServer() *Server {
	return &Server{
		statusMap: map[string]healthpb.HealthCheckResponse_ServingStatus{"": healthpb.HealthCheckResponse_SERVING},
		updates:   make(map[string]map[healthgrpc.Health_WatchServer]chan healthpb.HealthCheckResponse_ServingStatus),
	}
}

// Check implements `service Health`.
func (s *Server) Check(_ context.Context, in *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if servingStatus, ok := s.statusMap[in.Service]; ok {
		return &healthpb.HealthCheckResponse{
			Status: servingStatus,
		}, nil
	}
	return nil, status.Error(codes.NotFound, "unknown service")
}

// List implements `service Health`.
func (s *Server) List(_ context.Context, _ *healthpb.HealthListRequest) (*healthpb.HealthListResponse, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if len(s.statusMap) > maxAllowedServices {
		return nil, status.Errorf(codes.ResourceExhausted, "server health list exceeds maximum capacity: %d", maxAllowedServices)
	}

	statusMap := make(map[string]*healthpb.HealthCheckResponse, len(s.statusMap))
	for k, v := range s.statusMap {
		statusMap[k] = &healthpb.HealthCheckResponse{Status: v}
	}

	return &healthpb.HealthListResponse{Statuses: statusMap}, nil
}

// Watch implements `service Health`.
func (s *Server) Watch(in *healthpb.HealthCheckRequest, stream healthgrpc.Health_WatchServer) error {
	service := in.Service
	// update channel is used for getting service status updates.
	update := make(chan healthpb.HealthCheckResponse_ServingStatus, 1)
	s.mu.Lock()
	// Puts the initial status to the channel.
	if servingStatus, ok := s.statusMap[service]; ok {
		update <- servingStatus
	} else {
		update <- healthpb.HealthCheckResponse_SERVICE_UNKNOWN
	}

	// Registers the update channel to the correct place in the updates map.