[raft] client-side health check

enterprise/server/raft/cache/BUILD

@@ -9,6 +9,7 @@ go_library(
     deps = [
         "//enterprise/server/filestore",
         "//enterprise/server/raft/bringup",
+        "//enterprise/server/raft/constants",
         "//enterprise/server/raft/logger",
         "//enterprise/server/raft/rbuilder",
         "//enterprise/server/raft/registry",

enterprise/server/raft/cache/cache.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/filestore"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/bringup"
+	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/constants"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/rbuilder"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/registry"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/sender"
@@ -194,7 +195,7 @@ func Register(env *real_environment.RealEnv) error {
 			return nil
 		},
 	)
-	env.GetHealthChecker().AddHealthCheck("raft_cache", rc)
+	env.GetHealthChecker().AddHealthCheck(constants.HealthCheckName, rc)
 	return nil
 }
 

enterprise/server/raft/client/client.go

@@ -93,7 +93,7 @@ func (c *APIClient) getClient(ctx context.Context, peer string) (returnedClient
 		return rfspb.NewApiClient(conn), nil
 	}
 	log.Debugf("Creating new client for peer: %q", peer)
-	conn, err := grpc_client.DialSimple("grpc://" + peer)
+	conn, err := grpc_client.DialSimple("grpc://"+peer, grpc_client.HealthCheckGRPCOption(constants.HealthCheckName))
 	if err != nil {
 		return nil, err
 	}

enterprise/server/raft/constants/constants.go

@@ -19,7 +19,8 @@ const (
 	NodePartitionUsageEvent   = "node_partition_usage_event"
 	PlacementDriverQueryEvent = "placement_driver_query_event"
 
-	CacheName = "raft"
+	CacheName       = "raft"
+	HealthCheckName = "raft_cache"
 )
 
 // Key range contants

enterprise/server/raft/sender/BUILD

@@ -15,12 +15,14 @@ go_library(
         "//enterprise/server/raft/rbuilder",
         "//proto:raft_go_proto",
         "//proto:raft_service_go_proto",
+        "//server/metrics",
         "//server/util/log",
         "//server/util/proto",
         "//server/util/rangemap",
         "//server/util/retry",
         "//server/util/status",
         "//server/util/tracing",
+        "@com_github_prometheus_client_golang//prometheus",
         "@io_opentelemetry_go_otel//attribute",
         "@io_opentelemetry_go_otel//codes",
         "@org_golang_google_grpc//status",

enterprise/server/raft/sender/sender.go

@@ -11,12 +11,14 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/keys"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/rangecache"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/rbuilder"
+	"github.com/buildbuddy-io/buildbuddy/server/metrics"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/proto"
 	"github.com/buildbuddy-io/buildbuddy/server/util/rangemap"
 	"github.com/buildbuddy-io/buildbuddy/server/util/retry"
 	"github.com/buildbuddy-io/buildbuddy/server/util/status"
 	"github.com/buildbuddy-io/buildbuddy/server/util/tracing"
+	"github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
 
@@ -330,6 +332,9 @@ func (s *Sender) TryReplicas(ctx context.Context, rd *rfpb.RangeDescriptor, fn r
 			switch {
 			// range not found, no replicas are likely to have it; bail.
 			case strings.HasPrefix(m, constants.RangeNotFoundMsg), strings.HasPrefix(m, constants.RangeNotCurrentMsg):
+				metrics.RaftRangeNotPresentErrorCount.With(prometheus.Labels{
+					metrics.RaftNodeHostIDLabel: replica.GetNhid(),
+				}).Inc()
 				return 0, status.OutOfRangeErrorf("failed to TryReplicas on c%dn%d: no replicas are likely to have it: %s", replica.GetRangeId(), replica.GetReplicaId(), m)
 			case strings.HasPrefix(m, constants.RangeNotLeasedMsg), strings.HasPrefix(m, constants.RangeLeaseInvalidMsg):
 				logs = append(logs, fmt.Sprintf("skipping c%dn%d: out of range: %s", replica.GetRangeId(), replica.GetReplicaId(), err))

server/metrics/metrics.go

@@ -2384,6 +2384,15 @@ var (
 		Help:      "The total number of eviction errors",
 	})
 
+	RaftRangeNotPresentErrorCount = promauto.NewCounterVec(prometheus.CounterOpts{
+		Namespace: bbNamespace,
+		Subsystem: "raft",
+		Name:      "range_not_present_errors",
+		Help:      "The total number of range not present error",
+	}, []string{
+		RaftNodeHostIDLabel,
+	})
+
 	RaftListenerEventsDropped = promauto.NewCounterVec(prometheus.CounterOpts{
 		Namespace: bbNamespace,
 		Subsystem: "raft",

server/util/grpc_client/grpc_client.go

@@ -2,6 +2,7 @@ package grpc_client
 
 import (
 	"context"
+	"fmt"
 	"math"
 	"math/rand/v2"
 	"net/url"
@@ -36,6 +37,12 @@ const (
 
 var (
 	poolSize = flag.Int("grpc_client.pool_size", 15, "Number of connections to create to each target.")
+
+	serviceConfig = `{
+	"healthCheckConfig": {
+		"serviceName": "%s"
+	}
+}`
 )
 
 type clientConn struct {
@@ -337,3 +344,7 @@ func CommonGRPCClientOptions() []grpc.DialOption {
 		}),
 	}
 }
+
+func HealthCheckGRPCOption(serviceName string) grpc.DialOption {
+	return grpc.WithDefaultServiceConfig(fmt.Sprintf(serviceConfig, serviceName))
+}

server/util/healthcheck/healthcheck.go

@@ -32,8 +32,9 @@ var (
 )
 
 const (
-	healthCheckPeriod  = 3 * time.Second // The time to wait between health checks.
-	healthCheckTimeout = 2 * time.Second // How long a health check may take, max.
+	healthCheckPeriod        = 3 * time.Second // The time to wait between health checks.
+	healthCheckTimeout       = 2 * time.Second // How long a health check may take, max.
+	healthCheckWatchInterval = 5 * time.Second
 )
 
 type serviceStatus struct {
@@ -53,6 +54,8 @@ type HealthChecker struct {
 	shutdownFuncs []interfaces.CheckerFunc
 	readyToServe  bool
 	shuttingDown  bool
+
+	update chan struct{}
 }
 
 func NewHealthChecker(serverType string) *HealthChecker {
@@ -65,6 +68,7 @@ func NewHealthChecker(serverType string) *HealthChecker {
 		checkersMu:    sync.Mutex{},
 		checkers:      make(map[string]interfaces.Checker, 0),
 		lastStatus:    make([]*serviceStatus, 0),
+		update:        make(chan struct{}, 1),
 	}
 	sigTerm := make(chan os.Signal, 1)
 	go func() {
@@ -111,6 +115,7 @@ func (h *HealthChecker) handleShutdownFuncs() {
 	h.readyToServe = false
 	h.shuttingDown = true
 	h.mu.Unlock()
+	h.updateServingStatus()
 
 	// We use fmt here and below because this code is called from the
 	// signal handler and log.Printf can be a little wonky.
@@ -159,6 +164,7 @@ func (h *HealthChecker) AddHealthCheck(name string, f interfaces.Checker) {
 	h.mu.Lock()
 	h.readyToServe = false
 	h.mu.Unlock()
+	h.updateServingStatus()
 }
 
 func (h *HealthChecker) WaitForGracefulShutdown() {
@@ -226,6 +232,7 @@ func (h *HealthChecker) runHealthChecks(ctx context.Context) {
 		h.lastStatus = statusData
 	}
 	h.mu.Unlock()
+	h.updateServingStatus()
 
 	if newReadinessState != previousReadinessState {
 		log.Infof("HealthChecker transitioning from ready: %t => ready: %t", previousReadinessState, newReadinessState)
@@ -276,6 +283,34 @@ func (h *HealthChecker) LivenessHandler() http.Handler {
 	})
 }
 
+func (h *HealthChecker) servingStatus() hlpb.HealthCheckResponse_ServingStatus {
+	h.mu.RLock()
+	ready := h.readyToServe
+	shuttingDown := h.shuttingDown
+	h.mu.RUnlock()
+
+	if shuttingDown {
+		return hlpb.HealthCheckResponse_UNKNOWN
+	}
+
+	if ready {
+		return hlpb.HealthCheckResponse_SERVING
+	}
+	return hlpb.HealthCheckResponse_NOT_SERVING
+}
+
+func (h *HealthChecker) updateServingStatus() {
+	h.mu.Lock()
+
+	select {
+	case <-h.update:
+	default:
+	}
+	h.update <- struct{}{}
+
+	h.mu.Unlock()
+}
+
 func (h *HealthChecker) Check(ctx context.Context, req *hlpb.HealthCheckRequest) (*hlpb.HealthCheckResponse, error) {
 	// GRPC does not have indepenent health and readiness checks like HTTP does.
 	// An additional wrinkle is that AWS ALB's do not support sending a service
@@ -284,25 +319,29 @@ func (h *HealthChecker) Check(ctx context.Context, req *hlpb.HealthCheckRequest)
 	//   - SERVING when the service is ready
 	//   - NOT_SERVING when the service is not ready
 	//   - UNKNOWN when the service is shutting down.
-	h.mu.RLock()
-	ready := h.readyToServe
-	shuttingDown := h.shuttingDown
-	h.mu.RUnlock()
-	rsp := &hlpb.HealthCheckResponse{}
-	if ready {
-		rsp.Status = hlpb.HealthCheckResponse_SERVING
-	} else {
-		rsp.Status = hlpb.HealthCheckResponse_NOT_SERVING
-	}
-
-	if shuttingDown {
-		rsp.Status = hlpb.HealthCheckResponse_UNKNOWN
+	rsp := &hlpb.HealthCheckResponse{
+		Status: h.servingStatus(),
 	}
 	return rsp, nil
 }
 
 func (h *HealthChecker) Watch(req *hlpb.HealthCheckRequest, stream hlpb.Health_WatchServer) error {
-	return status.UnimplementedError("Watch not implemented")
+	currentStatus := hlpb.HealthCheckResponse_SERVICE_UNKNOWN
+
+	for {
+		select {
+		case <-stream.Context().Done():
+			return nil
+		case <-h.update:
+			newStatus := h.servingStatus()
+			if newStatus != currentStatus {
+				currentStatus = newStatus
+				if err := stream.Send(&hlpb.HealthCheckResponse{Status: currentStatus}); err != nil {
+					return err
+				}
+			}
+		}
+	}
 }
 
 func logGoroutineProfile() {