feat: skip memberlist zone-aware routing if there is a zone with members but no alive bridges (#800)

**What this PR does**:

When zone-aware routing is enabled, if there is a zone with members but
no alive bridges we end up in a network partitioning state. I think we
can make the routing logic more resilient by checking for this specific
condition, and internally disable the zone-aware routing if we detect a
zone without alive bridges.

The new logic is definitely slower than the previous one, but in
absolute terms I don't think it's prohibitive. `SelectNodes()` gets
called once every gossip interval (200ms) and takes 0.17ms with 10K
nodes and 2 zones (10K nodes is a pretty large scale):

```
goos: darwin
goarch: arm64
pkg: github.com/grafana/dskit/kv/memberlist
cpu: Apple M3 Pro
BenchmarkZoneAwareNodeSelectionDelegate_SelectNodes-11    	    7033	    170166 ns/op	   81922 B/op	       1 allocs/op
BenchmarkZoneAwareNodeSelectionDelegate_SelectNodes-11    	    6940	    170834 ns/op	   81922 B/op	       1 allocs/op
BenchmarkZoneAwareNodeSelectionDelegate_SelectNodes-11    	    6925	    171649 ns/op	   81922 B/op	       1 allocs/op
BenchmarkZoneAwareNodeSelectionDelegate_SelectNodes-11    	    6906	    213222 ns/op	   81922 B/op	       1 allocs/op
BenchmarkZoneAwareNodeSelectionDelegate_SelectNodes-11    	    6957	    172906 ns/op	   81922 B/op	       1 allocs/op
PASS
ok  	github.com/grafana/dskit/kv/memberlist	9.339s
```

This PR is based on grafana/memberlist#13.

**Which issue(s) this PR fixes**:

N/A

**Checklist**
- [x] Tests updated

---------

Signed-off-by: Marco Pracucci <marco@pracucci.com>

Author: pracucci

go.mod

@@ -136,4 +136,4 @@ require (
 
 // Replace memberlist with our fork which includes some fixes that haven't been
 // merged upstream yet.
-replace github.com/hashicorp/memberlist => github.com/grafana/memberlist v0.3.1-0.20251125151730-2125cd96c917
+replace github.com/hashicorp/memberlist => github.com/grafana/memberlist v0.3.1-0.20251126142931-6f9f62ab6f86

go.sum

@@ -139,8 +139,8 @@ github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
 github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
 github.com/grafana/gomemcache v0.0.0-20251008140118-65a671e12fdf h1:4fN03UTf6i7wZbLfxL+L/jAIiAxoThFi4rb9//SoGlY=
 github.com/grafana/gomemcache v0.0.0-20251008140118-65a671e12fdf/go.mod h1:j/s0jkda4UXTemDs7Pgw/vMT06alWc42CHisvYac0qw=
-github.com/grafana/memberlist v0.3.1-0.20251125151730-2125cd96c917 h1:9NzUQJ7ODlhiyYTc428KisLveJyA6Nd7rgqf/u6M/8U=
-github.com/grafana/memberlist v0.3.1-0.20251125151730-2125cd96c917/go.mod h1:h60o12SZn/ua/j0B6iKAZezA4eDaGsIuPO70eOaJ6WE=
+github.com/grafana/memberlist v0.3.1-0.20251126142931-6f9f62ab6f86 h1:aTwfQuroOmOr//QEn9J1MtC4R4CPR9/IbUd8hZrbWKo=
+github.com/grafana/memberlist v0.3.1-0.20251126142931-6f9f62ab6f86/go.mod h1:h60o12SZn/ua/j0B6iKAZezA4eDaGsIuPO70eOaJ6WE=
 github.com/grafana/otel-profiling-go v0.5.1 h1:stVPKAFZSa7eGiqbYuG25VcqYksR6iWvF3YH66t4qL8=
 github.com/grafana/otel-profiling-go v0.5.1/go.mod h1:ftN/t5A/4gQI19/8MoWurBEtC6gFw8Dns1sJZ9W4Tls=
 github.com/grafana/pyroscope-go/godeltaprof v0.1.9 h1:c1Us8i6eSmkW+Ez05d3co8kasnuOY813tbMN8i/a3Og=

kv/memberlist/memberlist_client.go

@@ -528,7 +528,7 @@ func (m *KV) configureZoneAwareRouting(mlCfg *memberlist.Config) error {
 	m.nodeMeta = localMeta
 
 	// Set up the node selection delegate.
-	mlCfg.NodeSelection = newZoneAwareNodeSelectionDelegate(role, m.cfg.ZoneAwareRouting.Zone, m.logger)
+	mlCfg.NodeSelection = newZoneAwareNodeSelectionDelegate(role, m.cfg.ZoneAwareRouting.Zone, m.logger, m.registerer)
 
 	// The bridge always prefer another bridge as first node. If the bridge only push/pull to 1 node per interval, then
 	// it will only communicate to bridges, potentially leading to network partitioning if the gossiping is not

kv/memberlist/node_zone_aware_routing.go

@@ -3,10 +3,14 @@ package memberlist
 import (
 	"flag"
 	"fmt"
+	"math/rand"
+	"slices"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/hashicorp/memberlist"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
 // NodeRole represents the role of a node in the memberlist cluster.
@@ -85,29 +89,107 @@ type zoneAwareNodeSelectionDelegate struct {
 	localRole NodeRole
 	localZone string
 	logger    log.Logger
+
+	// Metrics
+	selectNodesCalls        prometheus.Counter
+	selectNodesCallsSkipped prometheus.Counter
 }
 
 // newZoneAwareNodeSelectionDelegate creates a new zone-aware node selection delegate.
-func newZoneAwareNodeSelectionDelegate(localRole NodeRole, localZone string, logger log.Logger) *zoneAwareNodeSelectionDelegate {
+func newZoneAwareNodeSelectionDelegate(localRole NodeRole, localZone string, logger log.Logger, registerer prometheus.Registerer) *zoneAwareNodeSelectionDelegate {
 	return &zoneAwareNodeSelectionDelegate{
 		localRole: localRole,
 		localZone: localZone,
 		logger:    logger,
+		selectNodesCalls: promauto.With(registerer).NewCounter(prometheus.CounterOpts{
+			Name: "memberlist_client_zone_aware_routing_select_nodes_total",
+			Help: "Total number of times memberlist attempted to select node candidates for gossiping (tracked only when when zone-aware routing is enabled).",
+		}),
+		selectNodesCallsSkipped: promauto.With(registerer).NewCounter(prometheus.CounterOpts{
+			Name: "memberlist_client_zone_aware_routing_select_nodes_skipped_total",
+			Help: "Total number of times memberlist zone-aware routing was skipped because the local zone is unknown or a zone has no alive bridges.",
+		}),
 	}
 }
 
-// SelectNode implements memberlist.NodeSelectionDelegate.
-// It determines whether a remote node should be selected for gossip operations and whether it should be preferred.
-func (d *zoneAwareNodeSelectionDelegate) SelectNode(node memberlist.Node) (selected, preferred bool) {
-	// Decode remote node metadata.
-	remoteMeta := EncodedNodeMetadata(node.Meta)
-	remoteZone := remoteMeta.Zone()
-	remoteRole := remoteMeta.Role()
+// SelectNodes implements memberlist.NodeSelectionDelegate.
+// It determines which remote nodes should be selected for gossip operations and which one should be preferred.
+func (d *zoneAwareNodeSelectionDelegate) SelectNodes(nodes []*memberlist.NodeState) (selected []*memberlist.NodeState, preferred *memberlist.NodeState) {
+	d.selectNodesCalls.Inc()
+
+	if d.localRole != NodeRoleMember && d.localRole != NodeRoleBridge {
+		level.Warn(d.logger).Log("msg", "memberlist zone-aware routing is running with an unknown role", "role", d.localRole)
+	}
+
+	// Skip zone-aware routing if local zone is not set.
+	if d.localZone == "" {
+		d.selectNodesCallsSkipped.Inc()
+		return nodes, nil
+	}
+
+	// Pre-allocate backing arrays on the stack for up to 5 zones (common case).
+	zonesWithMembers := make([]string, 0, 5)
+	zonesWithAliveBridges := make([]string, 0, 5)
+
+	// Build selected slice and track zones in a single pass.
+	selected = make([]*memberlist.NodeState, 0, len(nodes))
+	preferredCount := 0 // Count of preferred candidates seen (for reservoir sampling).
+
+	for _, node := range nodes {
+		remoteMeta := EncodedNodeMetadata(node.Meta)
+		remoteZone := remoteMeta.Zone()
+		remoteRole := remoteMeta.Role()
+
+		// Track zones to check if any zone has members but no alive bridges.
+		if remoteZone != "" {
+			if remoteRole == NodeRoleBridge {
+				// Only count alive bridges.
+				if node.State == memberlist.StateAlive {
+					if !slices.Contains(zonesWithAliveBridges, remoteZone) {
+						zonesWithAliveBridges = append(zonesWithAliveBridges, remoteZone)
+					}
+				}
+			} else {
+				if !slices.Contains(zonesWithMembers, remoteZone) {
+					zonesWithMembers = append(zonesWithMembers, remoteZone)
+				}
+			}
+		}
+
+		// Apply zone-aware selection.
+		isSelected, isPreferred := d.selectNode(remoteZone, remoteRole)
+		if isSelected {
+			selected = append(selected, node)
+			if isPreferred {
+				preferredCount++
+				// Reservoir sampling: select this node with a probability of 1/preferredCount.
+				if rand.Intn(preferredCount) == 0 {
+					preferred = node
+				}
+			}
+		}
+	}
+
+	// Skip zone-aware routing if any zone has members but no alive bridges.
+	// This prevents network partitioning when bridges are missing or dead.
+	for _, zone := range zonesWithMembers {
+		if !slices.Contains(zonesWithAliveBridges, zone) {
+			d.selectNodesCallsSkipped.Inc()
+			level.Warn(d.logger).Log("msg", "memberlist zone-aware routing is skipped because a zone has no alive bridge", "zone", zone)
+			return nodes, nil
+		}
+	}
 
-	// If either the local zone or the remote zone are unknown, select the node but don't prefer it.
+	return selected, preferred
+}
+
+// selectNode determines whether a remote node should be selected for gossip operations
+// and whether it should be considered a preferred candidate.
+func (d *zoneAwareNodeSelectionDelegate) selectNode(remoteZone string, remoteRole NodeRole) (selected, preferredCandidate bool) {
+	// If the remote zone is unknown, select the node but don't prefer it.
 	// This prevents network partitioning: if every other memberlist node filters it out, then that
 	// remote node would not receive updates and would get isolated.
-	if d.localZone == "" || remoteZone == "" {
+	if remoteZone == "" {
 		return true, false
 	}
 
@@ -132,8 +214,6 @@ func (d *zoneAwareNodeSelectionDelegate) SelectNode(node memberlist.Node) (selec
 		return false, false
 
 	default:
-		level.Warn(d.logger).Log("msg", "memberlist zone-aware routing is running with an unknown role", "role", d.localRole)
-
 		// Unknown role: select but don't prefer (should never happen).
 		return true, false
 	}