From 37351c690cc47f1e9c95f2621b4a27c55c436665 Mon Sep 17 00:00:00 2001
From: Tony An <40644135+tonyjongyoonan@users.noreply.github.com>
Date: Thu, 13 Jul 2023 10:24:07 -0700
Subject: [PATCH] fixed StaticStrideScheduler looping too many times (#10370)

---
 .../xds/WeightedRoundRobinLoadBalancer.java   | 13 ++++++------
 .../WeightedRoundRobinLoadBalancerTest.java   | 20 +++++++++----------
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java b/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java
index d5d8c4d9e4..c2f06c8046 100644
--- a/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java
+++ b/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java
@@ -350,7 +350,6 @@ final class WeightedRoundRobinLoadBalancer extends RoundRobinLoadBalancer {
   @VisibleForTesting
   static final class StaticStrideScheduler {
     private final short[] scaledWeights;
-    private final int sizeDivisor;
     private final AtomicInteger sequence;
     private static final int K_MAX_WEIGHT = 0xFFFF;
 
@@ -373,7 +372,7 @@ final class WeightedRoundRobinLoadBalancer extends RoundRobinLoadBalancer {
       if (numWeightedChannels > 0) {
         meanWeight = (short) Math.round(scalingFactor * sumWeight / numWeightedChannels);
       } else {
-        meanWeight = 1;
+        meanWeight = (short) Math.round(scalingFactor);
       }
 
       // scales weights s.t. max(weights) == K_MAX_WEIGHT, meanWeight is scaled accordingly
@@ -387,7 +386,6 @@ final class WeightedRoundRobinLoadBalancer extends RoundRobinLoadBalancer {
       }
 
       this.scaledWeights = scaledWeights;
-      this.sizeDivisor = numChannels;
       this.sequence = new AtomicInteger(random.nextInt());
 
     }
@@ -433,15 +431,18 @@ final class WeightedRoundRobinLoadBalancer extends RoundRobinLoadBalancer {
      * an offset that varies per backend index is also included to the calculation.
      */
     int pick() {
+      int i = 0;
       while (true) {
+        i++;
         long sequence = this.nextSequence();
-        int backendIndex = (int) (sequence % this.sizeDivisor);
-        long generation = sequence / this.sizeDivisor;
-        int weight = Short.toUnsignedInt(this.scaledWeights[backendIndex]);
+        int backendIndex = (int) (sequence % scaledWeights.length);
+        long generation = sequence / scaledWeights.length;
+        int weight = Short.toUnsignedInt(scaledWeights[backendIndex]);
         long offset = (long) K_MAX_WEIGHT / 2 * backendIndex;
         if ((weight * generation + offset) % K_MAX_WEIGHT < K_MAX_WEIGHT - weight) {
           continue;
         }
+        assert i <= scaledWeights.length : "scheduler has more than one pass through";
         return backendIndex;
       }
     }
diff --git a/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java
index 879cac871b..505b2f7dff 100644
--- a/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java
+++ b/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java
@@ -330,11 +330,11 @@ public class WeightedRoundRobinLoadBalancerTest {
     }
     assertThat(pickCount.size()).isEqualTo(3);
     assertThat(Math.abs(pickCount.get(weightedSubchannel1) / 10000.0 - subchannel1PickRatio))
-        .isAtMost(0.001);
+        .isAtMost(0.0001);
     assertThat(Math.abs(pickCount.get(weightedSubchannel2) / 10000.0 - subchannel2PickRatio ))
-        .isAtMost(0.001);
+        .isAtMost(0.0001);
     assertThat(Math.abs(pickCount.get(weightedSubchannel3) / 10000.0 - subchannel3PickRatio ))
-        .isAtMost(0.001);
+        .isAtMost(0.0001);
   }
 
   @Test
@@ -751,12 +751,12 @@ public class WeightedRoundRobinLoadBalancerTest {
     }
     assertThat(pickCount.size()).isEqualTo(3);
     assertThat(Math.abs(pickCount.get(weightedSubchannel1) / 1000.0 - 4.0 / 9))
-            .isAtMost(0.002);
+            .isAtMost(0.001);
     assertThat(Math.abs(pickCount.get(weightedSubchannel2) / 1000.0 - 2.0 / 9))
-            .isAtMost(0.002);
+            .isAtMost(0.001);
     // subchannel3's weight is average of subchannel1 and subchannel2
     assertThat(Math.abs(pickCount.get(weightedSubchannel3) / 1000.0 - 3.0 / 9))
-            .isAtMost(0.002);
+            .isAtMost(0.001);
   }
 
   @Test
@@ -947,7 +947,7 @@ public class WeightedRoundRobinLoadBalancerTest {
     }
     for (int i = 0; i < 3; i++) {
       assertThat(Math.abs(pickCount.getOrDefault(i, 0) / 1000.0 - weights[i] / totalWeight))
-          .isAtMost(0.01);
+          .isAtMost(0.001);
     }
   }
 
@@ -964,7 +964,7 @@ public class WeightedRoundRobinLoadBalancerTest {
     }
     for (int i = 0; i < 3; i++) {
       assertThat(Math.abs(pickCount.getOrDefault(i, 0) / 1000.0 - weights[i] / totalWeight))
-          .isAtMost(0.01);
+          .isAtMost(0.001);
     }
   }
 
@@ -981,7 +981,7 @@ public class WeightedRoundRobinLoadBalancerTest {
     }
     for (int i = 0; i < 2; i++) {
       assertThat(Math.abs(pickCount.getOrDefault(i, 0) / 1000.0 - weights[i] / totalWeight))
-          .isAtMost(0.01);
+          .isAtMost(0.001);
     }
   }
 
@@ -1015,7 +1015,7 @@ public class WeightedRoundRobinLoadBalancerTest {
     }
     for (int i = 0; i < 8; i++) {
       assertThat(Math.abs(pickCount.getOrDefault(i, 0) / 1000.0 - weights[i] / totalWeight))
-          .isAtMost(0.01);
+          .isAtMost(0.004);
     }
   }