health: Add List method to gRPC Health service

health/server.go

@@ -30,6 +30,12 @@ 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 +68,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,6 +19,11 @@
 package health
 
 import (
+	"context"
+	"errors"
+	"fmt"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"sync"
 	"testing"
 	"time"
@@ -81,3 +86,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()
+	s.mu.Lock()
+	// Remove the zero value
+	delete(s.statusMap, "")
+	// Fill out status map with information
+	for i := 1; i <= 3; i++ {
+		s.statusMap[fmt.Sprintf("%d", i)] = healthpb.HealthCheckResponse_SERVING
+	}
+	s.mu.Unlock()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	var in healthpb.HealthListRequest
+	out, err := s.List(ctx, &in)
+
+	if err != nil {
+		t.Fatalf("s.List(ctx, &in) returned err %v, want nil", err)
+	}
+	if len(out.GetStatuses()) != len(s.statusMap) {
+		t.Fatalf("len(out.GetStatuses()) = %d, want %d", len(out.GetStatuses()), len(s.statusMap))
+	}
+	for key := range out.GetStatuses() {
+		v, ok := s.statusMap[key]
+		if !ok {
+			t.Fatalf("key %s does not exist in s.statusMap", key)
+		}
+		if v != healthpb.HealthCheckResponse_SERVING {
+			t.Fatalf("%s returned status %d, want %d", key, healthpb.HealthCheckResponse_SERVING, v)
+		}
+	}
+}
+
+// TestListResourceExhausted verifies that List() returns a ResourceExhausted error if no. of services are more than
+// maxAllowedServices.
+func (s) TestListResourceExhausted(t *testing.T) {
+	s := NewServer()
+	s.mu.Lock()
+	// Remove the zero value
+	delete(s.statusMap, "")
+
+	// Fill out status map with service information, 101 elements will trigger an error.
+	for i := 1; i <= maxAllowedServices+1; i++ {
+		s.statusMap[fmt.Sprintf("%d", i)] = healthpb.HealthCheckResponse_SERVING
+	}
+	s.mu.Unlock()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	var in healthpb.HealthListRequest
+	_, err := s.List(ctx, &in)
+
+	if err == nil {
+		t.Fatalf("s.List(ctx, &in) return nil error, want non-nil")
+	}
+	if !errors.Is(err, status.Errorf(codes.ResourceExhausted, "server health list exceeds maximum capacity: %d", maxAllowedServices)) {
+		t.Fatal("List should have failed with resource exhausted")
+	}
+}