health: Add List method to gRPC Health service (#8155)

Author: marcoshuck

health/server.go

@@ -30,6 +30,13 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+const (
+	// maxAllowedServices defines the maximum number of resources a List
+	// operation can return. An error is returned if the number of services
+	// exceeds this limit.
+	maxAllowedServices = 100
+)
+
 // Server implements `service Health`.
 type Server struct {
 	healthgrpc.UnimplementedHealthServer
@@ -62,6 +69,23 @@ func (s *Server) Check(_ context.Context, in *healthpb.HealthCheckRequest) (*hea
 	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

health/server_internal_test.go

@@ -19,10 +19,18 @@
 package health
 
 import (
+	"context"
+	"errors"
+	"fmt"
 	"sync"
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/testing/protocmp"
+
 	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/internal/grpctest"
 )
@@ -81,3 +89,65 @@ func (s) TestShutdown(t *testing.T) {
 		t.Fatalf("status for %s is %v, want %v", testService, status, healthpb.HealthCheckResponse_NOT_SERVING)
 	}
 }
+
+// TestList verifies that List() returns the health status of all the services if no. of services are within
+// maxAllowedLimits.
+func (s) TestList(t *testing.T) {
+	s := NewServer()
+
+	// Fill out status map with information
+	const length = 3
+	for i := 0; i < length; i++ {
+		s.SetServingStatus(fmt.Sprintf("%d", i),
+			healthpb.HealthCheckResponse_ServingStatus(i))
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	var in healthpb.HealthListRequest
+	got, err := s.List(ctx, &in)
+
+	if err != nil {
+		t.Fatalf("s.List(ctx, &in) returned err %v, want nil", err)
+	}
+	if len(got.GetStatuses()) != length+1 {
+		t.Fatalf("len(out.GetStatuses()) = %d, want %d",
+			len(got.GetStatuses()), length+1)
+	}
+	want := &healthpb.HealthListResponse{
+		Statuses: map[string]*healthpb.HealthCheckResponse{
+			"":  {Status: healthpb.HealthCheckResponse_SERVING},
+			"0": {Status: healthpb.HealthCheckResponse_UNKNOWN},
+			"1": {Status: healthpb.HealthCheckResponse_SERVING},
+			"2": {Status: healthpb.HealthCheckResponse_NOT_SERVING},
+		},
+	}
+	if diff := cmp.Diff(got, want, protocmp.Transform()); diff != "" {
+		t.Fatalf("Health response did not match expectation.  Diff (-got, +want): %s", diff)
+	}
+}
+
+// TestListResourceExhausted verifies that List()
+// returns a ResourceExhausted error if no. of services are more than
+// maxAllowedServices.
+func (s) TestListResourceExhausted(t *testing.T) {
+	s := NewServer()
+
+	// Fill out status map with service information,
+	// 101 (100 + 1 existing) elements will trigger an error.
+	for i := 1; i <= maxAllowedServices; i++ {
+		s.SetServingStatus(fmt.Sprintf("%d", i),
+			healthpb.HealthCheckResponse_SERVING)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	var in healthpb.HealthListRequest
+	_, err := s.List(ctx, &in)
+
+	want := status.Errorf(codes.ResourceExhausted,
+		"server health list exceeds maximum capacity: %d", maxAllowedServices)
+	if !errors.Is(err, want) {
+		t.Fatalf("s.List(ctx, &in) returned %v, want %v", err, want)
+	}
+}