balancer/ringhash: fix uint32 weight sum overflow in normalizeWeights

balancer/ringhash/ring.go

@@ -137,12 +137,19 @@ func newRing(endpoints *resolver.EndpointMap[*endpointState], minRingSize, maxRi
 //
 // Must be called with a non-empty endpoints map.
 func normalizeWeights(endpoints *resolver.EndpointMap[*endpointState]) ([]endpointInfo, float64) {
-	var weightSum uint32
+	// Accumulate in a uint64 so the sum cannot wrap: each weight is a uint32
+	// and control-plane supplied localities/endpoints can make the total exceed
+	// math.MaxUint32. A wrapped sum is smaller than the real one, so the
+	// normalized weights come out greater than 1 and the ring grows past its
+	// configured max size. For example, a wrapped uint32 sum can land on zero,
+	// which would turn the division below into +Inf and make newRing spin
+	// forever building the ring.
+	var weightSum uint64
 	// Since attributes are explicitly ignored in the EndpointMap key, we need
 	// to iterate over the values to get the weights.
 	endpointVals := endpoints.Values()
 	for _, epState := range endpointVals {
-		weightSum += epState.weight
+		weightSum += uint64(epState.weight)
 	}
 	ret := make([]endpointInfo, 0, endpoints.Len())
 	min := 1.0

balancer/ringhash/ring_test.go

@@ -76,6 +76,74 @@ func (s) TestRingNew(t *testing.T) {
 	}
 }
 
+// TestRingNewWeightSumOverflow checks that endpoint weights whose sum exceeds
+// math.MaxUint32 do not wrap the weight accumulator to zero. A zero sum used to
+// make normalizeWeights divide by zero, producing +Inf normalized weights and
+// an unbounded ring-build loop in newRing.
+func (s) TestRingNewWeightSumOverflow(t *testing.T) {
+	endpoints := []resolver.Endpoint{
+		testEndpoint("a", 1<<31),
+		testEndpoint("b", 1<<31), // sum is exactly 2^32, wraps a uint32 to 0.
+	}
+	m := resolver.NewEndpointMap[*endpointState]()
+	m.Set(endpoints[0], &endpointState{hashKey: "a", weight: 1 << 31})
+	m.Set(endpoints[1], &endpointState{hashKey: "b", weight: 1 << 31})
+
+	const min, max uint64 = 1024, 4096
+	r := newRing(m, min, max, nil)
+	if got := uint64(len(r.items)); got < min || got > max {
+		t.Fatalf("newRing built a ring of size %d, want within [%d, %d]", got, min, max)
+	}
+	for _, e := range endpoints {
+		var count int
+		for _, ii := range r.items {
+			if ii.hashKey == hashKey(e) {
+				count++
+			}
+		}
+		got := float64(count) / float64(len(r.items))
+		if got != 0.5 {
+			t.Fatalf("endpoint %q occupies %v of the ring, want 0.5", hashKey(e), got)
+		}
+	}
+}
+
+// TestRingNewWeightSumOverflowToNonZero checks that endpoint weights whose sum
+// exceeds math.MaxUint32 and wraps to a non-zero value still produce a ring
+// within the configured size bounds. A wrapped non-zero sum used to make
+// normalizeWeights return weights greater than 1, growing the ring past
+// maxRingSize.
+func (s) TestRingNewWeightSumOverflowToNonZero(t *testing.T) {
+	endpoints := []resolver.Endpoint{
+		testEndpoint("a", 1<<31),
+		testEndpoint("b", 1<<31),
+		testEndpoint("c", 1<<30), // sum is 2^32 + 2^30, wraps a uint32 to 2^30.
+	}
+	m := resolver.NewEndpointMap[*endpointState]()
+	m.Set(endpoints[0], &endpointState{hashKey: "a", weight: 1 << 31})
+	m.Set(endpoints[1], &endpointState{hashKey: "b", weight: 1 << 31})
+	m.Set(endpoints[2], &endpointState{hashKey: "c", weight: 1 << 30})
+
+	const min, max uint64 = 1024, 4096
+	r := newRing(m, min, max, nil)
+	if got := uint64(len(r.items)); got < min || got > max {
+		t.Fatalf("newRing built a ring of size %d, want within [%d, %d]", got, min, max)
+	}
+	wantFractions := map[string]float64{"a": 0.4, "b": 0.4, "c": 0.2}
+	for _, e := range endpoints {
+		var count int
+		for _, ii := range r.items {
+			if ii.hashKey == hashKey(e) {
+				count++
+			}
+		}
+		got := float64(count) / float64(len(r.items))
+		if want := wantFractions[hashKey(e)]; !equalApproximately(got, want) {
+			t.Fatalf("endpoint %q occupies %v of the ring, want ~%v", hashKey(e), got, want)
+		}
+	}
+}
+
 func equalApproximately(x, y float64) bool {
 	delta := math.Abs(x - y)
 	mean := math.Abs(x+y) / 2.0