client_routes: add sticky route tracking

When a host has multiple routes (one per connection), remember which
connectionID was used on the first successful TranslateHost call and
prefer it on subsequent calls via findPreferredRoute. If the preferred
route is removed (e.g. connection pruned), fall back to FindByHostID.

This avoids unnecessary connection churn when multiple PrivateLink
endpoints serve the same host.

Author: sylwiaszunejko

client_routes.go

@@ -187,15 +187,16 @@ func (l *UnresolvedClientRouteList) FindByHostID(hostID string) *UnresolvedClien
 }
 
 type ClientRoutesHandler struct {
-	log         StdLogger
-	c           controlConnection
-	resolver    DNSResolver
-	sub         *eventbus.Subscriber[events.Event]
-	routes      UnresolvedClientRouteList
-	updateTasks chan updateTask
-	closeChan   chan struct{}
-	cfg         ClientRoutesConfig
-	mu          sync.RWMutex
+	log          StdLogger
+	c            controlConnection
+	resolver     DNSResolver
+	sub          *eventbus.Subscriber[events.Event]
+	routes       UnresolvedClientRouteList
+	stickyRoute  map[string]string // hostID → preferred connectionID
+	updateTasks  chan updateTask
+	closeChan    chan struct{}
+	cfg          ClientRoutesConfig
+	mu           sync.RWMutex
 	pickTLSPorts bool
 	initialized  bool
 }
@@ -215,6 +216,20 @@ func pickProperPort(pickTLSPorts bool, rec *UnresolvedClientRoute) uint16 {
 	return rec.CQLPort
 }
 
+// findPreferredRoute returns the route for hostID that matches the sticky
+// connectionID, falling back to the first route for that host.
+// Must be called with p.mu held (at least RLock).
+func (p *ClientRoutesHandler) findPreferredRoute(hostID string) *UnresolvedClientRoute {
+	if preferred, ok := p.stickyRoute[hostID]; ok {
+		for i := range p.routes {
+			if p.routes[i].HostID == hostID && p.routes[i].ConnectionID == preferred {
+				return &p.routes[i]
+			}
+		}
+	}
+	return p.routes.FindByHostID(hostID)
+}
+
 // TranslateHost implements AddressTranslatorV2 interface.
 // It resolves DNS on every call rather than caching resolved addresses.
 func (p *ClientRoutesHandler) TranslateHost(host AddressTranslatorHostInfo, addr AddressPort) (AddressPort, error) {
@@ -224,7 +239,7 @@ func (p *ClientRoutesHandler) TranslateHost(host AddressTranslatorHostInfo, addr
 	}
 
 	p.mu.RLock()
-	rec := p.routes.FindByHostID(hostID)
+	rec := p.findPreferredRoute(hostID)
 	var route UnresolvedClientRoute
 	found := rec != nil
 	if found {
@@ -249,6 +264,10 @@ func (p *ClientRoutesHandler) TranslateHost(host AddressTranslatorHostInfo, addr
 		return addr, fmt.Errorf("record %s/%s has target port empty", route.HostID, route.ConnectionID)
 	}
 
+	p.mu.Lock()
+	p.stickyRoute[hostID] = route.ConnectionID
+	p.mu.Unlock()
+
 	return AddressPort{Address: ips[0], Port: port}, nil
 }
 
@@ -400,6 +419,7 @@ func NewClientRoutesAddressTranslator(
 		updateTasks:  make(chan updateTask, 1024),
 		resolver:     resolver,
 		routes:       make(UnresolvedClientRouteList, 0),
+		stickyRoute:  make(map[string]string),
 	}
 }
 

client_routes_unit_test.go

@@ -119,8 +119,9 @@ func TestClientRoutesHandlerTranslateHost(t *testing.T) {
 	})
 
 	handler := &ClientRoutesHandler{
-		resolver: resolver,
-		routes:   make(UnresolvedClientRouteList, 0),
+		stickyRoute: make(map[string]string),
+		resolver:    resolver,
+		routes:      make(UnresolvedClientRouteList, 0),
 	}
 
 	res, err := handler.TranslateHost(noHost, addr)
@@ -162,14 +163,68 @@ func TestClientRoutesHandlerTranslateHost(t *testing.T) {
 		resolver: dnsResolverFunc(func(host string) ([]net.IP, error) {
 			return nil, errors.New("lookup failed")
 		}),
-		routes: UnresolvedClientRouteList{{ConnectionID: "c2", HostID: "h2", Address: "host", CQLPort: 9042}},
+		stickyRoute: make(map[string]string),
+		routes:      UnresolvedClientRouteList{{ConnectionID: "c2", HostID: "h2", Address: "host", CQLPort: 9042}},
 	}
 	_, err = errorHandler.TranslateHost(testHostInfo{hostID: "h2"}, addr)
 	if err == nil {
 		t.Fatalf("expected resolver error to bubble up")
 	}
 }
 
+func TestTranslateHost_StickyRoute(t *testing.T) {
+	addr := AddressPort{Address: net.ParseIP("1.1.1.1"), Port: 9042}
+	resolvedIPs := map[string]net.IP{
+		"addr-c1": net.ParseIP("10.0.0.1"),
+		"addr-c2": net.ParseIP("10.0.0.2"),
+	}
+	handler := &ClientRoutesHandler{
+		pickTLSPorts: false,
+		stickyRoute:  make(map[string]string),
+		resolver: dnsResolverFunc(func(host string) ([]net.IP, error) {
+			if ip, ok := resolvedIPs[host]; ok {
+				return []net.IP{ip}, nil
+			}
+			return nil, fmt.Errorf("unknown host %s", host)
+		}),
+		routes: UnresolvedClientRouteList{
+			{ConnectionID: "c1", HostID: "h1", Address: "addr-c1", CQLPort: 9042},
+			{ConnectionID: "c2", HostID: "h1", Address: "addr-c2", CQLPort: 9042},
+		},
+	}
+
+	// First call picks the first route (c1) and records it as sticky.
+	res, err := handler.TranslateHost(testHostInfo{hostID: "h1"}, addr)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !res.Address.Equal(net.ParseIP("10.0.0.1")) {
+		t.Fatalf("expected first route IP 10.0.0.1, got %v", res.Address)
+	}
+
+	// Second call should stick to c1 even though c2 also matches h1.
+	res, err = handler.TranslateHost(testHostInfo{hostID: "h1"}, addr)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !res.Address.Equal(net.ParseIP("10.0.0.1")) {
+		t.Fatalf("expected sticky route IP 10.0.0.1, got %v", res.Address)
+	}
+
+	// Remove c1 route; sticky route should fall back to c2.
+	handler.mu.Lock()
+	handler.routes = handler.routes[1:]
+	handler.mu.Unlock()
+
+	res, err = handler.TranslateHost(testHostInfo{hostID: "h1"}, addr)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !res.Address.Equal(net.ParseIP("10.0.0.2")) {
+		t.Fatalf("expected fallback to c2 IP 10.0.0.2, got %v", res.Address)
+	}
+}
+
 func TestGetHostPortMappingFromClusterQuery(t *testing.T) {
 	tcases := []struct {
 		name          string