balancer/ringhash: address review comments

Check exact 0.5 distribution in overflow test, add a test for a sum
that wraps to a non-zero value, and expand the accumulator comment to
cover both wrap outcomes.

Author: nvxbug

balancer/ringhash/ring.go

@@ -139,8 +139,11 @@ func newRing(endpoints *resolver.EndpointMap[*endpointState], minRingSize, maxRi
 func normalizeWeights(endpoints *resolver.EndpointMap[*endpointState]) ([]endpointInfo, float64) {
 	// 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 uint32 sum can land on zero, which would turn the
-	// division below into +Inf and make newRing spin forever building the ring.
+	// 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.

balancer/ringhash/ring_test.go

@@ -102,8 +102,44 @@ func (s) TestRingNewWeightSumOverflow(t *testing.T) {
 			}
 		}
 		got := float64(count) / float64(len(r.items))
-		if !equalApproximately(got, 0.5) {
-			t.Fatalf("endpoint %q occupies %v of the ring, want ~0.5", hashKey(e), got)
+		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)
 		}
 	}
 }