fix: resolve PR comments

epam · Jan 31, 2025 · 5ef1f93 · 5ef1f93
1 parent f7a330b
commit 5ef1f93
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 13 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
@@ -54,19 +54,14 @@ public Upstream next() {
             return availableUpstreams.get(0);
         }
         int total = availableUpstreams.stream().map(Upstream::getWeight).reduce(0, Integer::sum);
-        // make sure the upper bound `total` is inclusive
-        // the lowest bound should be 1 otherwise the 1st upstream with weight 1 has more possibility
-        // to be selected because zero is included in its range
-        // e.g. there are 3 upstreams with equal weight = 1 and the ranges are [0,1], [2,2] and [3,3]
-        // definitely the 1st upstream has higher possibility
-        int random = generator.nextInt(1, 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
@@ -31,25 +31,25 @@ void testWeightedLoadBalancer() {
 
         RandomizedWeightedBalancer balancer = new RandomizedWeightedBalancer("model1", upstreams, generator);
 
-        when(generator.nextInt(1, 11)).thenReturn(0);
+        when(generator.nextInt(10)).thenReturn(0);
 
         Upstream upstream = balancer.next();
         assertNotNull(upstream);
         assertEquals(upstreams.get(0), upstream);
 
-        when(generator.nextInt(1, 11)).thenReturn(2);
+        when(generator.nextInt(10)).thenReturn(2);
 
         upstream = balancer.next();
         assertNotNull(upstream);
         assertEquals(upstreams.get(1), upstream);
 
-        when(generator.nextInt(1, 11)).thenReturn(6);
+        when(generator.nextInt(10)).thenReturn(5);
 
         upstream = balancer.next();
         assertNotNull(upstream);
         assertEquals(upstreams.get(2), upstream);
 
-        when(generator.nextInt(1, 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(1, 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);