Merge pull request Netflix#365 from Aloren/bugfix/weighted-response-t…

…ime-rule-index-out-of-bound Fix IndexOutOfBoundsException for WeightedResponseTimeRule
wolongqin · Jan 31, 2018 · 80f33f3 · 80f33f3
2 parents 015cc51 + 26f1926
commit 80f33f3
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 2 deletions.
diff --git a/ribbon-loadbalancer/build.gradle b/ribbon-loadbalancer/build.gradle
@@ -8,6 +8,8 @@ dependencies {
     compile "com.netflix.archaius:archaius-core:${archaius_version}"
     compile 'com.netflix.netflix-commons:netflix-commons-util:0.1.1'
     testCompile 'junit:junit:4.11'
+    testCompile 'org.mockito:mockito-core:2.13.0'
+    testCompile 'org.awaitility:awaitility:3.0.0'
     testCompile 'org.slf4j:slf4j-log4j12:1.7.2'
     testCompile "com.sun.jersey:jersey-server:${jersey_version}"
 }
diff --git a/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/WeightedResponseTimeRule.java b/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/WeightedResponseTimeRule.java
@@ -24,6 +24,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Random;
 import java.util.Timer;
@@ -150,6 +151,10 @@ public void shutdown() {
         }
     }
 
+    List<Double> getAccumulatedWeights() {
+        return Collections.unmodifiableList(accumulatedWeights);
+    }
+
     @edu.umd.cs.findbugs.annotations.SuppressWarnings(value = "RCN_REDUNDANT_NULLCHECK_OF_NULL_VALUE")
     @Override
     public Server choose(ILoadBalancer lb, Object key) {
@@ -178,7 +183,7 @@ public Server choose(ILoadBalancer lb, Object key) {
             double maxTotalWeight = currentWeights.size() == 0 ? 0 : currentWeights.get(currentWeights.size() - 1); 
             // No server has been hit yet and total weight is not initialized
             // fallback to use round robin
-            if (maxTotalWeight < 0.001d) {
+            if (maxTotalWeight < 0.001d || serverCount != currentWeights.size()) {
                 server =  super.choose(getLoadBalancer(), key);
                 if(server == null) {
                     return server;
@@ -279,7 +284,7 @@ public void maintainWeights() {
     void setWeights(List<Double> weights) {
         this.accumulatedWeights = weights;
     }
-    
+
     @Override
     public void initWithNiwsConfig(IClientConfig clientConfig) {
         super.initWithNiwsConfig(clientConfig);

diff --git a/ribbon-loadbalancer/src/test/java/com/netflix/loadbalancer/WeightedResponseTimeRuleTest.java b/ribbon-loadbalancer/src/test/java/com/netflix/loadbalancer/WeightedResponseTimeRuleTest.java
@@ -0,0 +1,91 @@
+package com.netflix.loadbalancer;
+
+import org.awaitility.core.ThrowingRunnable;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.List;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singletonList;
+import static org.awaitility.Awaitility.await;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class WeightedResponseTimeRuleTest {
+
+    private static final Object KEY = "key";
+
+    private AbstractLoadBalancer loadBalancer;
+    private WeightedResponseTimeRule rule;
+
+    @Before
+    public void setUp() throws Exception {
+        rule = new WeightedResponseTimeRule();
+        loadBalancer = mock(AbstractLoadBalancer.class);
+        setupLoadBalancer(asList(server("first"), server("second"), server("third")));
+        rule.setLoadBalancer(loadBalancer);
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        rule.shutdown();
+    }
+
+    @Test
+    public void shouldNotFailWithIndexOutOfBoundExceptionWhenChoosingServerWhenNumberOfServersIsDecreased() throws Exception {
+        waitUntilWeightsAreCalculated();
+
+        setupLoadBalancer(singletonList(server("other")));
+
+        Server chosen = rule.choose(loadBalancer, KEY);
+
+        assertNotNull(chosen);
+    }
+
+    private void waitUntilWeightsAreCalculated() {
+        await().untilAsserted(new ThrowingRunnable() {
+            @Override
+            public void run() throws Throwable {
+                List<Double> weights = rule.getAccumulatedWeights();
+                assertNotEquals(weights.size(), 0);
+            }
+        });
+    }
+
+    private AbstractLoadBalancer setupLoadBalancer(List<Server> servers) {
+        LoadBalancerStats loadBalancerStats = getLoadBalancerStats(servers);
+        when(loadBalancer.getLoadBalancerStats()).thenReturn(loadBalancerStats);
+        when(loadBalancer.getReachableServers()).thenReturn(servers);
+        when(loadBalancer.getAllServers()).thenReturn(servers);
+        return loadBalancer;
+    }
+
+    private LoadBalancerStats getLoadBalancerStats(List<Server> servers) {
+        LoadBalancerStats stats = mock(LoadBalancerStats.class);
+        // initialize first server with maximum response time
+        // so that we could reproduce issue with decreased number of servers in loadbalancer
+        int responseTimeMax = servers.size() * 100;
+        for (Server server : servers) {
+            ServerStats s1 = statsWithResponseTimeAverage(responseTimeMax);
+            when(stats.getSingleServerStat(server)).thenReturn(s1);
+            responseTimeMax -= 100;
+        }
+        return stats;
+    }
+
+    private ServerStats statsWithResponseTimeAverage(double responseTimeAverage) {
+        ServerStats serverStats = mock(ServerStats.class);
+        when(serverStats.getResponseTimeAvg()).thenReturn(responseTimeAverage);
+        return serverStats;
+    }
+
+    private Server server(String id) {
+        Server server = new Server(id);
+        server.setAlive(true);
+        return server;
+    }
+}