Use kube-resolver to resolve addresses in raft registry (#11422)

1. Created a podWatcherManager that manages podWatchers. This is
extracted from
kubeResolverBuilder, so that the code can be re-used.
2. Created WatchPodIP() API for non-grpc consumers
3. Resolve() now returns resolved IP if available. The connection key is
the
   resolved IP instead of the hostname; so that when the IP changes, new
requests get pushed to a new queue instead of the old queue with the old
IP
   in dragonboat.

Author: luluz66

enterprise/server/raft/client/BUILD

@@ -15,6 +15,7 @@ go_library(
         "//server/environment",
         "//server/metrics",
         "//server/util/grpc_client",
+        "//server/util/kuberesolver",
         "//server/util/lockmap",
         "//server/util/log",
         "//server/util/proto",

enterprise/server/raft/client/client.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"flag"
 	"fmt"
-	"os"
 	"strconv"
 	"sync"
 	"time"
@@ -16,6 +15,7 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/server/environment"
 	"github.com/buildbuddy-io/buildbuddy/server/metrics"
 	"github.com/buildbuddy-io/buildbuddy/server/util/grpc_client"
+	"github.com/buildbuddy-io/buildbuddy/server/util/kuberesolver"
 	"github.com/buildbuddy-io/buildbuddy/server/util/lockmap"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/proto"
@@ -104,7 +104,7 @@ func (c *APIClient) getClient(ctx context.Context, peer string) (returnedClient
 	// Use kube:/// resolver when running in k8s for instant IP updates on
 	// pod restarts. Fall back to default grpc:// resolver otherwise (e.g. tests).
 	target := "grpc://" + peer
-	if os.Getenv("KUBERNETES_SERVICE_HOST") != "" {
+	if kuberesolver.RunningInKubernetes() {
 		target = "kube:///" + peer
 	}
 	conn, err := grpc_client.DialSimple(target, grpc.WithConnectParams(grpc.ConnectParams{

enterprise/server/raft/registry/BUILD

@@ -10,6 +10,7 @@ go_library(
         "//enterprise/server/raft/constants",
         "//proto:raft_go_proto",
         "//server/interfaces",
+        "//server/util/kuberesolver",
         "//server/util/log",
         "//server/util/proto",
         "//server/util/status",
@@ -28,8 +29,12 @@ go_test(
         ":registry",
         "//server/gossip",
         "//server/testutil/testport",
+        "//server/util/kuberesolver",
         "//server/util/log",
         "//server/util/status",
         "@com_github_stretchr_testify//require",
+        "@io_k8s_api//core/v1:core",
+        "@io_k8s_apimachinery//pkg/apis/meta/v1:meta",
+        "@io_k8s_client_go//kubernetes/fake",
     ],
 )

enterprise/server/raft/registry/registry.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/constants"
 	"github.com/buildbuddy-io/buildbuddy/server/interfaces"
+	"github.com/buildbuddy-io/buildbuddy/server/util/kuberesolver"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/proto"
 	"github.com/buildbuddy-io/buildbuddy/server/util/status"
@@ -68,6 +69,13 @@ type StaticRegistry struct {
 
 	targetAddresses sync.Map // map of NHID(string) => addresses
 
+	// Pod watcher for resolving raft addresses to IPs via k8s Watch API.
+	// When set, raft addresses that are pod FQDNs are watched for IP changes
+	// and Resolve() returns the resolved IP instead of the hostname.
+	podWatcherManager *kuberesolver.PodWatcherManager
+	resolvedRaftAddrs sync.Map // map of raftAddress(string) => resolved IP:port(string)
+	raftWatchCancels  sync.Map // map of raftAddress(string) => cancel func()
+
 	log log.Logger
 }
 
@@ -85,6 +93,52 @@ func NewStaticNodeRegistry(streamConnections uint64, v dbConfig.TargetValidator,
 	return n
 }
 
+// SetPodWatcherManager configures the registry to resolve raft addresses
+// (pod FQDNs) to IPs using the k8s Watch API via the given PodWatcherManager.
+// When set, Resolve() returns the resolved pod IP instead of the hostname,
+// and uses the IP as the connection key so that an IP change triggers a new
+// connection.
+func (n *StaticRegistry) SetPodWatcherManager(m *kuberesolver.PodWatcherManager) {
+	n.podWatcherManager = m
+}
+
+// resolveRaftAddress returns the resolved IP:port for the given raft address
+// if a pod watcher is tracking it, otherwise returns the address as-is.
+func (n *StaticRegistry) resolveRaftAddress(raftAddr string) string {
+	if resolved, ok := n.resolvedRaftAddrs.Load(raftAddr); ok {
+		return resolved.(string)
+	}
+	return raftAddr
+}
+
+// watchRaftAddress starts watching the given raft address for IP changes
+// via the k8s Watch API. If the address is not a pod FQDN or the pod watcher
+// is not configured, this is a no-op.
+func (n *StaticRegistry) watchRaftAddress(raftAddr string) {
+	if n.podWatcherManager == nil {
+		return
+	}
+	if _, ok := n.raftWatchCancels.Load(raftAddr); ok {
+		return
+	}
+	cancel, err := n.podWatcherManager.WatchPodIP(raftAddr, func(ipPort string, watchErr error) {
+		if watchErr != nil {
+			if _, had := n.resolvedRaftAddrs.LoadAndDelete(raftAddr); had {
+				n.log.Warningf("Raft address %s lost resolution: %s", raftAddr, watchErr)
+			}
+			return
+		}
+		if prev, loaded := n.resolvedRaftAddrs.Swap(raftAddr, ipPort); !loaded || prev.(string) != ipPort {
+			n.log.Infof("Raft address %s resolved to %s", raftAddr, ipPort)
+		}
+	})
+	if err != nil {
+		// Not a pod FQDN or not in k8s — fall back to hostname resolution.
+		return
+	}
+	n.raftWatchCancels.Store(raftAddr, cancel)
+}
+
 // Add adds the specified node and its target info to the registry.
 func (n *StaticRegistry) Add(rangeID uint64, replicaID uint64, target string) {
 	if n.validate != nil && !n.validate(target) {
@@ -156,7 +210,8 @@ func (n *StaticRegistry) Resolve(rangeID uint64, replicaID uint64) (string, stri
 	if err != nil {
 		return "", "", err
 	}
-	return ci.GetRaftAddress(), n.getConnectionKey(ci.GetRaftAddress(), rangeID), nil
+	addr := n.resolveRaftAddress(ci.GetRaftAddress())
+	return addr, n.getConnectionKey(addr, rangeID), nil
 }
 
 // ResolveRaft returns the raft address and the connection key of the specified node.
@@ -224,6 +279,7 @@ func (n *StaticRegistry) AddNode(target, raftAddress, grpcAddress string) {
 		grpc: grpcAddress,
 	}
 	n.targetAddresses.Store(target, a)
+	n.watchRaftAddress(raftAddress)
 }
 
 // ListNodes lists all the {NHID, raftAddress, grpcAddress} available in the
@@ -244,6 +300,12 @@ func (n *StaticRegistry) ListNodes() []*rfpb.ConnectionInfo {
 }
 
 func (n *StaticRegistry) Close() error {
+	n.raftWatchCancels.Range(func(key, value interface{}) bool {
+		if cancel, ok := value.(func()); ok && cancel != nil {
+			cancel()
+		}
+		return true
+	})
 	return nil
 }
 
@@ -301,6 +363,12 @@ func NewDynamicNodeRegistry(gossipManager interfaces.GossipService, streamConnec
 	return dnr
 }
 
+// SetPodWatcherManager configures the underlying static registry to resolve
+// raft addresses via the k8s Watch API.
+func (d *DynamicNodeRegistry) SetPodWatcherManager(m *kuberesolver.PodWatcherManager) {
+	d.sReg.SetPodWatcherManager(m)
+}
+
 func (d *DynamicNodeRegistry) handleEvent(event *serf.UserEvent) {
 	if event.Name != constants.RegistryUpdateEvent {
 		return
@@ -462,7 +530,8 @@ func (d *DynamicNodeRegistry) Resolve(rangeID uint64, replicaID uint64) (string,
 	if err != nil {
 		return "", "", err
 	}
-	return ci.GetRaftAddress(), d.sReg.getConnectionKey(ci.GetRaftAddress(), rangeID), nil
+	addr := d.sReg.resolveRaftAddress(ci.GetRaftAddress())
+	return addr, d.sReg.getConnectionKey(addr, rangeID), nil
 }
 
 // Lookup returns the connectionInfo of the specified node

enterprise/server/raft/registry/registry_test.go

@@ -4,13 +4,19 @@ import (
 	"context"
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/raft/registry"
 	"github.com/buildbuddy-io/buildbuddy/server/gossip"
 	"github.com/buildbuddy-io/buildbuddy/server/testutil/testport"
+	"github.com/buildbuddy-io/buildbuddy/server/util/kuberesolver"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/status"
 	"github.com/stretchr/testify/require"
+	"k8s.io/client-go/kubernetes/fake"
+
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 func init() {
@@ -143,6 +149,72 @@ func TestDynamicRegistryRemoveShard(t *testing.T) {
 	require.True(t, status.IsNotFoundError(err))
 }
 
+func TestStaticRegistryResolveWithPodWatcher(t *testing.T) {
+	pod := &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "server-0",
+			Namespace: "ns",
+		},
+		Status: corev1.PodStatus{
+			PodIP: "10.0.0.1",
+		},
+	}
+	k8sClient := fake.NewClientset(pod)
+	m := kuberesolver.NewPodWatcherManager(k8sClient)
+
+	nr := registry.NewStaticNodeRegistry(1, nil, log.Logger{})
+	nr.SetPodWatcherManager(m)
+
+	raftAddr := "server-0.headless.ns.svc.cluster.local:7238"
+	nr.Add(1, 1, "nhid-1")
+	nr.AddNode("nhid-1", raftAddr, "grpcaddress:1")
+
+	// Resolve should return the resolved pod IP instead of the hostname.
+	require.Eventually(t, func() bool {
+		addr, _, err := nr.Resolve(1, 1)
+		return err == nil && addr == "10.0.0.1:7238"
+	}, 5*time.Second, 10*time.Millisecond, "expected Resolve to return resolved IP")
+
+	// Connection key should use the resolved IP, not the hostname.
+	addr, key1, err := nr.Resolve(1, 1)
+	require.NoError(t, err)
+	require.Equal(t, "10.0.0.1:7238", addr)
+	require.Contains(t, key1, "10.0.0.1:7238")
+	require.NotContains(t, key1, "server-0.headless")
+
+	// ResolveGRPC is unaffected by pod watcher (resolves by NHID).
+	grpcAddr, err := nr.ResolveGRPC(context.Background(), "nhid-1")
+	require.NoError(t, err)
+	require.Equal(t, "grpcaddress:1", grpcAddr)
+
+	// Wait for the watch to be established before updating.
+	require.Eventually(t, func() bool {
+		for _, a := range k8sClient.Actions() {
+			if a.GetVerb() == "watch" {
+				return true
+			}
+		}
+		return false
+	}, 5*time.Second, 10*time.Millisecond, "timed out waiting for watch")
+
+	// Simulate a pod restart with a new IP.
+	pod.Status.PodIP = "10.0.0.2"
+	_, err = k8sClient.CoreV1().Pods("ns").Update(
+		context.Background(), pod, metav1.UpdateOptions{},
+	)
+	require.NoError(t, err)
+
+	// Resolve should return the new IP and a different connection key.
+	require.Eventually(t, func() bool {
+		addr, _, err := nr.Resolve(1, 1)
+		return err == nil && addr == "10.0.0.2:7238"
+	}, 5*time.Second, 10*time.Millisecond, "expected Resolve to return updated IP")
+
+	_, key2, err := nr.Resolve(1, 1)
+	require.NoError(t, err)
+	require.NotEqual(t, key1, key2, "connection key should change when IP changes")
+}
+
 func TestDynamicRegistryResolution(t *testing.T) {
 	node1Addr := localAddr(t)
 	node2Addr := localAddr(t)

enterprise/server/raft/store/BUILD

@@ -38,6 +38,7 @@ go_library(
         "//server/util/canary",
         "//server/util/disk",
         "//server/util/grpc_server",
+        "//server/util/kuberesolver",
         "//server/util/lib/set",
         "//server/util/log",
         "//server/util/proto",

enterprise/server/raft/store/store.go

@@ -49,6 +49,7 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/server/util/canary"
 	"github.com/buildbuddy-io/buildbuddy/server/util/disk"
 	"github.com/buildbuddy-io/buildbuddy/server/util/grpc_server"
+	"github.com/buildbuddy-io/buildbuddy/server/util/kuberesolver"
 	"github.com/buildbuddy-io/buildbuddy/server/util/lib/set"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/proto"
@@ -181,6 +182,9 @@ type registryHolder struct {
 func (rc *registryHolder) Create(nhid string, streamConnections uint64, v dbConfig.TargetValidator) (raftio.INodeRegistry, error) {
 	nhLog := log.NamedSubLogger(nhid)
 	r := registry.NewDynamicNodeRegistry(rc.g, streamConnections, v, nhLog)
+	if kuberesolver.RunningInKubernetes() {
+		r.SetPodWatcherManager(kuberesolver.DefaultManager())
+	}
 	rc.r = r
 	r.AddNode(nhid, rc.raftAddr, rc.grpcAddr)
 	return r, nil