[raft] client-side health check

buildbuddy-io · Feb 27, 2025 · 4e6d2ea · 4e6d2ea
1 parent a2c326c
commit 4e6d2ea
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 3 deletions.
diff --git a/enterprise/server/raft/cache/BUILD b/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",

diff --git a/enterprise/server/raft/cache/cache.go b/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
 }
 

diff --git a/enterprise/server/raft/client/client.go b/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
 	}

diff --git a/enterprise/server/raft/constants/constants.go b/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

diff --git a/server/util/grpc_client/grpc_client.go b/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))
+}