client_routes: add scoped pruning to Merge for stale entry removal

Add scopeConnectionIDs and scopeHostIDs parameters to Merge so it can
prune stale entries before upserting:
- Both non-empty (partial update): prune entries matching BOTH lists
- Only scopeConnectionIDs (full refresh): prune all entries for those connections
- Both empty: additive-only merge, no pruning

updateHostPortMapping now passes connection/host IDs through to Merge,
enabling automatic cleanup of decommissioned hosts.

Author: sylwiaszunejko

client_routes.go

@@ -135,9 +135,24 @@ func (l *UnresolvedClientRouteList) Len() int {
 }
 
 // Merge upserts entries from incoming into the list.
-// Existing entries matching by (ConnectionID, HostID) are updated in place;
-// new entries are appended.
-func (l *UnresolvedClientRouteList) Merge(incoming UnresolvedClientRouteList) {
+// Before upserting it prunes stale entries within the query scope defined by
+// scopeConnectionIDs and scopeHostIDs:
+//   - Both non-empty (partial update): prune entries matching BOTH lists.
+//   - Only scopeConnectionIDs (full refresh): prune all entries for those connections.
+//   - Both empty / nil: additive-only merge, no pruning.
+func (l *UnresolvedClientRouteList) Merge(incoming UnresolvedClientRouteList, scopeConnectionIDs, scopeHostIDs []string) {
+	switch {
+	case len(scopeConnectionIDs) > 0 && len(scopeHostIDs) > 0:
+		*l = slices.DeleteFunc(*l, func(r UnresolvedClientRoute) bool {
+			return slices.Contains(scopeConnectionIDs, r.ConnectionID) &&
+				slices.Contains(scopeHostIDs, r.HostID)
+		})
+	case len(scopeConnectionIDs) > 0:
+		*l = slices.DeleteFunc(*l, func(r UnresolvedClientRoute) bool {
+			return slices.Contains(scopeConnectionIDs, r.ConnectionID)
+		})
+	}
+
 	for _, inc := range incoming {
 		found := false
 		for id, existing := range *l {
@@ -355,7 +370,7 @@ func (p *ClientRoutesHandler) updateHostPortMapping(connectionIDs []string, host
 	}
 
 	p.mu.Lock()
-	p.routes.Merge(incoming)
+	p.routes.Merge(incoming, connectionIDs, hostIDs)
 	p.mu.Unlock()
 
 	return nil

client_routes_unit_test.go

@@ -66,15 +66,15 @@ func TestMerge(t *testing.T) {
 	// Same record: no change expected.
 	list.Merge(UnresolvedClientRouteList{
 		{ConnectionID: "c1", HostID: "h1", Address: "a1", CQLPort: 9042},
-	})
+	}, nil, nil)
 	if len(list) != 1 {
 		t.Fatalf("expected 1 entry, got %d", len(list))
 	}
 
 	// Updated address: record should be replaced.
 	list.Merge(UnresolvedClientRouteList{
 		{ConnectionID: "c1", HostID: "h1", Address: "a2", CQLPort: 9043},
-	})
+	}, nil, nil)
 	if list[0].Address != "a2" || list[0].CQLPort != 9043 {
 		t.Fatalf("expected record to update")
 	}
@@ -83,7 +83,7 @@ func TestMerge(t *testing.T) {
 	list = UnresolvedClientRouteList{}
 	list.Merge(UnresolvedClientRouteList{
 		{ConnectionID: "c2", HostID: "h2", Address: "a3", CQLPort: 9044},
-	})
+	}, nil, nil)
 	if len(list) != 1 {
 		t.Fatalf("expected new record to be appended")
 	}
@@ -226,6 +226,99 @@ func TestGetHostPortMappingFromClusterQuery(t *testing.T) {
 	}
 }
 
+func TestMerge_DeletedHost(t *testing.T) {
+	list := UnresolvedClientRouteList{
+		{ConnectionID: "c1", HostID: "h1", Address: "a1", CQLPort: 9042},
+		{ConnectionID: "c1", HostID: "h2", Address: "a2", CQLPort: 9042},
+	}
+
+	// Simulate event for (c1, h1) where query returned nothing → (c1,h1) should be removed.
+	list.Merge(nil, []string{"c1"}, []string{"h1"})
+
+	if len(list) != 1 {
+		t.Fatalf("expected 1 entry after pruning deleted host, got %d", len(list))
+	}
+	if list[0].HostID != "h2" {
+		t.Fatalf("expected h2 to survive, got %s", list[0].HostID)
+	}
+}
+
+func TestMerge_UpdatedHost(t *testing.T) {
+	list := UnresolvedClientRouteList{
+		{ConnectionID: "c1", HostID: "h1", Address: "old-addr", CQLPort: 9042},
+		{ConnectionID: "c2", HostID: "h1", Address: "old-addr2", CQLPort: 9042},
+		{ConnectionID: "c1", HostID: "h2", Address: "keep", CQLPort: 9042},
+	}
+
+	// h1 address changed; fresh query returns new data for h1. h2 is not affected.
+	list.Merge(UnresolvedClientRouteList{
+		{ConnectionID: "c1", HostID: "h1", Address: "new-addr", CQLPort: 9043},
+		{ConnectionID: "c2", HostID: "h1", Address: "new-addr2", CQLPort: 9043},
+	}, []string{"c1", "c2"}, []string{"h1"})
+
+	if len(list) != 3 {
+		t.Fatalf("expected 3 entries, got %d", len(list))
+	}
+	for _, r := range list {
+		if r.HostID == "h1" {
+			if r.Address != "new-addr" && r.Address != "new-addr2" {
+				t.Fatalf("expected h1 entries to have new addresses, got %s", r.Address)
+			}
+		}
+		if r.HostID == "h2" && r.Address != "keep" {
+			t.Fatalf("expected h2 entry to be preserved unchanged")
+		}
+	}
+}
+
+func TestMerge_FullRefresh_PrunesAllStale(t *testing.T) {
+	list := UnresolvedClientRouteList{
+		{ConnectionID: "c1", HostID: "h1", Address: "a1", CQLPort: 9042},
+		{ConnectionID: "c1", HostID: "h2", Address: "a2", CQLPort: 9042},
+		{ConnectionID: "c1", HostID: "h3", Address: "a3", CQLPort: 9042},
+	}
+
+	// Full refresh for connection c1: all entries for c1 are pruned, only h1 and h2 returned.
+	list.Merge(UnresolvedClientRouteList{
+		{ConnectionID: "c1", HostID: "h1", Address: "a1", CQLPort: 9042},
+		{ConnectionID: "c1", HostID: "h2", Address: "a2", CQLPort: 9042},
+	}, []string{"c1"}, nil)
+
+	if len(list) != 2 {
+		t.Fatalf("expected 2 entries after full refresh prune, got %d", len(list))
+	}
+	for _, r := range list {
+		if r.HostID == "h3" {
+			t.Fatalf("expected h3 to be pruned")
+		}
+	}
+}
+
 // TestUpdateHostPortMapping_FullRefresh_PrunesStaleEntries simulates the same
 // sequence of operations that updateHostPortMapping performs (lock → Merge → unlock)
 // to verify that a full refresh correctly prunes a host that disappeared.
+func TestUpdateHostPortMapping_FullRefresh_PrunesStaleEntries(t *testing.T) {
+	// Existing routes: h1, h2, h3.
+	routes := UnresolvedClientRouteList{
+		{ConnectionID: "c1", HostID: "h1", Address: "a1", CQLPort: 9042},
+		{ConnectionID: "c1", HostID: "h2", Address: "a2", CQLPort: 9042},
+		{ConnectionID: "c1", HostID: "h3", Address: "a3", CQLPort: 9042},
+	}
+
+	// Cluster now returns only h1 and h2 (h3 was decommissioned).
+	incoming := UnresolvedClientRouteList{
+		{ConnectionID: "c1", HostID: "h1", Address: "a1", CQLPort: 9042},
+		{ConnectionID: "c1", HostID: "h2", Address: "a2", CQLPort: 9042},
+	}
+
+	routes.Merge(incoming, []string{"c1"}, nil)
+
+	if len(routes) != 2 {
+		t.Fatalf("expected 2 entries after full-refresh prune, got %d", len(routes))
+	}
+	for _, r := range routes {
+		if r.HostID == "h3" {
+			t.Fatalf("h3 should have been pruned by full refresh")
+		}
+	}
+}