fix: RandomizedWeightedBalancer chooses the first upstream with highe…

…r probability than others #666 (#667)
epam · Feb 3, 2025 · b6d07a2 · b6d07a2
1 parent d8432f8
commit b6d07a2
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 10 deletions.
diff --git a/server/src/main/java/com/epam/aidial/core/server/upstream/RandomizedWeightedBalancer.java b/server/src/main/java/com/epam/aidial/core/server/upstream/RandomizedWeightedBalancer.java
@@ -50,16 +50,18 @@ public Upstream next() {
         if (availableUpstreams.isEmpty()) {
             return null;
         }
+        if (availableUpstreams.size() == 1) {
+            return availableUpstreams.get(0);
+        }
         int total = availableUpstreams.stream().map(Upstream::getWeight).reduce(0, Integer::sum);
-        // make sure the upper bound `total` is inclusive
-        int random = generator.nextInt(total + 1);
+        int random = generator.nextInt(total);
         int current = 0;
 
         Upstream result = null;
 
         for (Upstream upstream : availableUpstreams) {
             current += upstream.getWeight();
-            if (current >= random) {
+            if (current > random) {
                 result = upstream;
                 break;
             }

diff --git a/...er/src/test/java/com/epam/aidial/core/server/upstream/RandomizedWeightedBalancerTest.java b/...er/src/test/java/com/epam/aidial/core/server/upstream/RandomizedWeightedBalancerTest.java
@@ -19,7 +19,7 @@ public class RandomizedWeightedBalancerTest {
 
     @Mock
     private Random generator;
-    
+
     @Test
     void testWeightedLoadBalancer() {
         List<Upstream> upstreams = List.of(
@@ -31,25 +31,25 @@ void testWeightedLoadBalancer() {
 
         RandomizedWeightedBalancer balancer = new RandomizedWeightedBalancer("model1", upstreams, generator);
 
-        when(generator.nextInt(11)).thenReturn(0);
+        when(generator.nextInt(10)).thenReturn(0);
 
         Upstream upstream = balancer.next();
         assertNotNull(upstream);
         assertEquals(upstreams.get(0), upstream);
 
-        when(generator.nextInt(11)).thenReturn(2);
+        when(generator.nextInt(10)).thenReturn(2);
 
         upstream = balancer.next();
         assertNotNull(upstream);
         assertEquals(upstreams.get(1), upstream);
 
-        when(generator.nextInt(11)).thenReturn(6);
+        when(generator.nextInt(10)).thenReturn(5);
 
         upstream = balancer.next();
         assertNotNull(upstream);
         assertEquals(upstreams.get(2), upstream);
 
-        when(generator.nextInt(11)).thenReturn(10);
+        when(generator.nextInt(10)).thenReturn(9);
 
         upstream = balancer.next();
         assertNotNull(upstream);

diff --git a/server/src/test/java/com/epam/aidial/core/server/upstream/TieredBalancerTest.java b/server/src/test/java/com/epam/aidial/core/server/upstream/TieredBalancerTest.java
@@ -110,8 +110,8 @@ void testUpstreamFallback() {
                 .map(index  -> new Upstream("endpoint" + index, null, null, 1, 1))
                 .toList();
         model.setUpstreams(upstreams);
-        AtomicInteger counter = new AtomicInteger();
-        when(generator.nextInt(5)).thenAnswer(cb -> counter.incrementAndGet());
+        AtomicInteger counter = new AtomicInteger(-1);
+        when(generator.nextInt(4)).thenAnswer(cb -> counter.incrementAndGet());
         Supplier<Random> factory = () -> generator;
 
         UpstreamRouteProvider upstreamRouteProvider = new UpstreamRouteProvider(vertx, factory);