ringhash: fix normalizeWeights

xds/internal/balancer/ringhash/ring.go

@@ -122,24 +122,20 @@ func newRing(subConns *resolver.AddressMap, minRingSize, maxRingSize uint64, log
 // Must be called with a non-empty subConns map.
 func normalizeWeights(subConns *resolver.AddressMap) ([]subConnWithWeight, float64) {
 	var weightSum uint32
-	keys := subConns.Keys()
-	for _, a := range keys {
-		weightSum += getWeightAttribute(a)
+	for _, a := range subConns.Values() {
+		weightSum += a.(*subConn).weight
 	}
-	ret := make([]subConnWithWeight, 0, len(keys))
+	ret := make([]subConnWithWeight, 0, len(subConns.Keys()))
 	min := float64(1.0)
-	for _, a := range keys {
-		v, _ := subConns.Get(a)
-		scInfo := v.(*subConn)
+	for _, a := range subConns.Values() {
+		scInfo := a.(*subConn)
 		// getWeightAttribute() returns 1 if the weight attribute is not found
 		// on the address. And since this function is guaranteed to be called
 		// with a non-empty subConns map, weightSum is guaranteed to be
 		// non-zero. So, we need not worry about divide a by zero error here.
-		nw := float64(getWeightAttribute(a)) / float64(weightSum)
+		nw := float64(scInfo.weight) / float64(weightSum)
 		ret = append(ret, subConnWithWeight{sc: scInfo, weight: nw})
-		if nw < min {
-			min = nw
-		}
+		min = math.Min(min, nw)
 	}
 	// Sort the addresses to return consistent results.
 	//

xds/internal/balancer/ringhash/ringhash_test.go

@@ -344,17 +344,26 @@ func (s) TestThreeSubConnsAffinityMultiple(t *testing.T) {
 	}
 }
 
+// TestAddrWeightChange covers the following scenarios after setting up the
+// balancer with 3 addresses [A, B, C]:
+//   - updates balancer with [A, B, C], a new Picker should not be sent.
+//   - updates balancer with [A, B] (C removed), a new Picker is sent and the
+//     ring is updated.
+//   - updates balancer with [A, B], but B has a weight of 2, a new Picker is
+//     sent.  And the new ring should contain the correct number of entries
+//     and weights.
 func (s) TestAddrWeightChange(t *testing.T) {
-	wantAddrs := []resolver.Address{
+	addrs := []resolver.Address{
 		{Addr: testBackendAddrStrs[0]},
 		{Addr: testBackendAddrStrs[1]},
 		{Addr: testBackendAddrStrs[2]},
 	}
-	cc, b, p0 := setupTest(t, wantAddrs)
+	cc, b, p0 := setupTest(t, addrs)
 	ring0 := p0.(*picker).ring
 
+	// Update with the same addresses, should not send a new Picker.
 	if err := b.UpdateClientConnState(balancer.ClientConnState{
-		ResolverState:  resolver.State{Addresses: wantAddrs},
+		ResolverState:  resolver.State{Addresses: addrs},
 		BalancerConfig: testConfig,
 	}); err != nil {
 		t.Fatalf("UpdateClientConnState returned err: %v", err)
@@ -407,6 +416,23 @@ func (s) TestAddrWeightChange(t *testing.T) {
 	if p2.(*picker).ring == ring1 {
 		t.Fatalf("new picker after changing address weight has the same ring as before, want different")
 	}
+	// Ideally with the new update, the ring should contain 3 entries. 1 for the
+	// first address, and 2 for the second address.
+	if len(p2.(*picker).ring.items) != 3 {
+		t.Fatalf("new picker after changing address weight has %d entries, want 3", len(p2.(*picker).ring.items))
+	}
+	for _, i := range p2.(*picker).ring.items {
+		if i.sc.addr == testBackendAddrStrs[0] {
+			if i.sc.weight != 1 {
+				t.Fatalf("new picker after changing address weight has weight %d for %v, want 1", i.sc.weight, i.sc.addr)
+			}
+		}
+		if i.sc.addr == testBackendAddrStrs[1] {
+			if i.sc.weight != 2 {
+				t.Fatalf("new picker after changing address weight has weight %d for %v, want 2", i.sc.weight, i.sc.addr)
+			}
+		}
+	}
 }
 
 // TestSubConnToConnectWhenOverallTransientFailure covers the situation when the