fix: Place only READY SubConns on the hashring so RPCs are not queued behind a connecting or failed backend

.github/workflows/lint-test.yaml

@@ -8,7 +8,7 @@ on:  # yamllint disable-line rule:truthy
   pull_request:
     branches: ["*"]
 env:
-  GO_VERSION: "~1.23.0"
+  GO_VERSION: "~1.25.0"
 jobs:
   go-lint:
     name: "Lint Go"

balancer.go

@@ -146,13 +146,15 @@ func (b *builder) Name() string { return BalancerName }
 
 func (b *builder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer {
 	bal := &ringBalancer{
-		cc:       cc,
-		subConns: resolver.NewAddressMap(),
-		scStates: make(map[balancer.SubConn]connectivity.State),
-		csEvltr:  &balancer.ConnectivityStateEvaluator{},
-		state:    connectivity.Connecting,
-		hasher:   b.hashfn,
-		picker:   base.NewErrPicker(balancer.ErrNoSubConnAvailable),
+		cc:          cc,
+		subConns:    resolver.NewAddressMapV2[any](),
+		scStates:    make(map[balancer.SubConn]connectivity.State),
+		scKeys:      make(map[balancer.SubConn]string),
+		ringMembers: make(map[balancer.SubConn]struct{}),
+		csEvltr:     &balancer.ConnectivityStateEvaluator{},
+		state:       connectivity.Connecting,
+		hasher:      b.hashfn,
+		picker:      base.NewErrPicker(balancer.ErrNoSubConnAvailable),
 	}
 
 	return bal
@@ -186,8 +188,12 @@ type ringBalancer struct {
 	cc       balancer.ClientConn
 	picker   balancer.Picker
 	csEvltr  *balancer.ConnectivityStateEvaluator
-	subConns *resolver.AddressMap
+	subConns *resolver.AddressMapV2[any]
 	scStates map[balancer.SubConn]connectivity.State
+	// scKeys holds the hashring key of every SubConn the resolver gave us.
+	scKeys map[balancer.SubConn]string
+	// ringMembers is the set of SubConns on the hashring: only READY ones.
+	ringMembers map[balancer.SubConn]struct{}
 
 	config   *BalancerConfig
 	hashring *hashring.Ring
@@ -235,8 +241,18 @@ func (b *ringBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 		svcConfig := s.BalancerConfig.(*BalancerConfig)
 		if b.config == nil || svcConfig.ReplicationFactor != b.config.ReplicationFactor {
 			b.hashring = hashring.MustNew(b.hasher, svcConfig.ReplicationFactor)
-			b.config = svcConfig
+			// The new ring starts empty: put every READY SubConn back on it.
+			b.ringMembers = make(map[balancer.SubConn]struct{})
+			for sc, st := range b.scStates {
+				if st != connectivity.Ready {
+					continue
+				}
+				if err := b.addToRing(sc); err != nil {
+					return err
+				}
+			}
 		}
+		b.config = svcConfig
 	}
 
 	// if there's no hashring yet, the balancer hasn't yet parsed an initial
@@ -252,12 +268,16 @@ func (b *ringBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 	// if any new targets have been added, they are added to the hashring, and
 	// any that have been removed since the last update are removed from the
 	// hashring.
-	addrsSet := resolver.NewAddressMap()
+	addrsSet := resolver.NewAddressMapV2[any]()
 	for _, addr := range s.ResolverState.Addresses {
 		addrsSet.Set(addr, nil)
 
 		if _, ok := b.subConns.Get(addr); !ok {
 			// addr is addr new address (not existing in b.subConns).
+			// NewSubConn is deprecated only to warn that a SubConn will soon
+			// hold a single address. We already pass exactly one address, and
+			// grpc offers no replacement API yet.
+			//nolint:staticcheck // SA1019: single-address usage is the future-proof form
 			sc, err := b.cc.NewSubConn([]resolver.Address{addr}, balancer.NewSubConnOptions{HealthCheckEnabled: false})
 			if err != nil {
 				logger.Warningf("base.baseBalancer: failed to create new SubConn: %v", err)
@@ -266,15 +286,9 @@ func (b *ringBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 
 			b.subConns.Set(addr, sc)
 			b.scStates[sc] = connectivity.Idle
+			b.scKeys[sc] = addr.ServerName + addr.Addr
 			b.csEvltr.RecordTransition(connectivity.Shutdown, connectivity.Idle)
 			sc.Connect()
-
-			if err := b.hashring.Add(subConnMember{
-				SubConn: sc,
-				key:     addr.ServerName + addr.Addr,
-			}); err != nil {
-				return fmt.Errorf("couldn't add to hashring")
-			}
 		}
 	}
 
@@ -283,15 +297,12 @@ func (b *ringBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 		sc := sci.(balancer.SubConn)
 		// addr was removed by resolver.
 		if _, ok := addrsSet.Get(addr); !ok {
-			b.cc.RemoveSubConn(sc)
+			sc.Shutdown()
 			b.subConns.Delete(addr)
 			// Keep the state of this sc in b.scStates until sc's state becomes Shutdown.
 			// The entry will be deleted in UpdateSubConnState.
-			if err := b.hashring.Remove(subConnMember{
-				SubConn: sc,
-				key:     addr.ServerName + addr.Addr,
-			}); err != nil {
-				return fmt.Errorf("couldn't add to hashring")
+			if err := b.removeFromRing(sc); err != nil {
+				return err
 			}
 		}
 	}
@@ -313,19 +324,52 @@ func (b *ringBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 		return balancer.ErrBadResolverState
 	}
 
-	// If the overall connection state is not in transient failure, we return
-	// addr new picker with addr reference to the hashring (otherwise an error picker)
+	b.regeneratePicker()
+	b.cc.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
+
+	return nil
+}
+
+// regeneratePicker installs an error picker while the balancer is in
+// TRANSIENT_FAILURE and a hashring picker otherwise.
+func (b *ringBalancer) regeneratePicker() {
 	if b.state == connectivity.TransientFailure {
 		b.picker = base.NewErrPicker(errors.Join(b.connErr, b.resolverErr))
-	} else {
-		b.picker = &picker{
-			hashring: b.hashring,
-			spread:   b.config.Spread,
-		}
+		return
 	}
 
-	// update the ClientConn with the current hashring picker picker
-	b.cc.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
+	b.picker = &picker{
+		hashring: b.hashring,
+		spread:   b.config.Spread,
+	}
+}
+
+// addToRing places sc on the hashring if it is not already there.
+func (b *ringBalancer) addToRing(sc balancer.SubConn) error {
+	if _, ok := b.ringMembers[sc]; ok {
+		return nil
+	}
+
+	if err := b.hashring.Add(subConnMember{SubConn: sc, key: b.scKeys[sc]}); err != nil {
+		return fmt.Errorf("couldn't add to hashring: %w", err)
+	}
+
+	b.ringMembers[sc] = struct{}{}
+
+	return nil
+}
+
+// removeFromRing takes sc off the hashring if it is there.
+func (b *ringBalancer) removeFromRing(sc balancer.SubConn) error {
+	if _, ok := b.ringMembers[sc]; !ok {
+		return nil
+	}
+
+	if err := b.hashring.Remove(subConnMember{SubConn: sc, key: b.scKeys[sc]}); err != nil {
+		return fmt.Errorf("couldn't remove from hashring: %w", err)
+	}
+
+	delete(b.ringMembers, sc)
 
 	return nil
 }
@@ -362,19 +406,34 @@ func (b *ringBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.Su
 
 	b.scStates[sc] = s
 
+	// A backend only receives traffic while its connection is READY. Keys
+	// hashed to a backend that is connecting or failed move to the closest
+	// ready member instead of waiting on a connection that may never come.
+	var ringErr error
+	if s == connectivity.Ready {
+		ringErr = b.addToRing(sc)
+	} else {
+		ringErr = b.removeFromRing(sc)
+	}
+	if ringErr != nil {
+		logger.Warningf("consistent-hashring: %v", ringErr)
+	}
+
 	switch s {
 	case connectivity.Idle:
 		sc.Connect()
 	case connectivity.Shutdown:
 		// When an address was removed by resolver, b called RemoveSubConn but
 		// kept the sc's state in scStates. Remove state for this sc here.
 		delete(b.scStates, sc)
+		delete(b.scKeys, sc)
 	case connectivity.TransientFailure:
 		// Save error to be reported via picker.
 		b.connErr = state.ConnectionError
 	}
 
 	b.state = b.csEvltr.RecordTransition(oldS, s)
+	b.regeneratePicker()
 
 	b.cc.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
 }
@@ -401,25 +460,31 @@ var _ balancer.Picker = (*picker)(nil)
 // The value stored in CtxKey is hashed into the hashring, and the resulting
 // subconnection is used.
 //
-// There is no fallback behavior if the subconnection is unavailable; this
-// prevents the request from going to a node that doesn't expect to receive it.
-// As long as you are using a resolver that removes connections from the list
-// when they are observably unavailable, this is a non-issue.
+// The hashring only contains READY subconnections, so a key whose closest
+// backend is down is served by the next closest ready backend. When no
+// backend is ready the RPC is queued until one becomes ready (or the
+// balancer reports TRANSIENT_FAILURE and installs an error picker).
 //
 // Spread can be increased to be robust against single node availability
 // problems. If spread is greater than 1, a random selection is made from the
 // set of subconns matching the hash.
 func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
 	key := info.Ctx.Value(CtxKey).([]byte)
 
+	// FindN only fails with hashring.ErrNotEnoughMembers.
 	members, err := p.hashring.FindN(key, p.spread)
 	if err != nil {
-		return balancer.PickResult{}, err
+		// Fewer ready backends than the configured spread: use those that are.
+		members, err = p.hashring.FindN(key, 1)
+		if err != nil {
+			// No ready backends at all: queue the RPC until one is ready.
+			return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+		}
 	}
 
 	index := 0
-	if p.spread > 1 {
-		index = intn(p.spread)
+	if len(members) > 1 {
+		index = intn(uint8(len(members)))
 	}
 
 	chosen := members[index].(subConnMember)