Optimize leastActiveSelect and weight test case

dubbo-cluster/src/main/java/org/apache/dubbo/rpc/cluster/loadbalance/LeastActiveLoadBalance.java

@@ -16,7 +16,6 @@
  */
 package org.apache.dubbo.rpc.cluster.loadbalance;
 
-import org.apache.dubbo.common.Constants;
 import org.apache.dubbo.common.URL;
 import org.apache.dubbo.rpc.Invocation;
 import org.apache.dubbo.rpc.Invoker;
@@ -27,7 +26,6 @@
 
 /**
  * LeastActiveLoadBalance
- *
  */
 public class LeastActiveLoadBalance extends AbstractLoadBalance {
 
@@ -39,26 +37,26 @@ protected <T> Invoker<T> doSelect(List<Invoker<T>> invokers, URL url, Invocation
         int leastActive = -1; // The least active value of all invokers
         int leastCount = 0; // The number of invokers having the same least active value (leastActive)
         int[] leastIndexes = new int[length]; // The index of invokers having the same least active value (leastActive)
-        int totalWeight = 0; // The sum of weights
+        int totalWeight = 0; // The sum of with warmup weights
         int firstWeight = 0; // Initial value, used for comparision
         boolean sameWeight = true; // Every invoker has the same weight value?
         for (int i = 0; i < length; i++) {
             Invoker<T> invoker = invokers.get(i);
             int active = RpcStatus.getStatus(invoker.getUrl(), invocation.getMethodName()).getActive(); // Active number
-            int weight = invoker.getUrl().getMethodParameter(invocation.getMethodName(), Constants.WEIGHT_KEY, Constants.DEFAULT_WEIGHT); // Weight
+            int afterWarmup = getWeight(invoker, invocation);
             if (leastActive == -1 || active < leastActive) { // Restart, when find a invoker having smaller least active value.
                 leastActive = active; // Record the current least active value
                 leastCount = 1; // Reset leastCount, count again based on current leastCount
                 leastIndexes[0] = i; // Reset
-                totalWeight = weight; // Reset
-                firstWeight = weight; // Record the weight the first invoker
+                totalWeight = afterWarmup; // Reset
+                firstWeight = afterWarmup; // Record the weight the first invoker
                 sameWeight = true; // Reset, every invoker has the same weight value?
             } else if (active == leastActive) { // If current invoker's active value equals with leaseActive, then accumulating.
                 leastIndexes[leastCount++] = i; // Record index number of this invoker
-                totalWeight += weight; // Add this invoker's weight to totalWeight.
+                totalWeight += afterWarmup; // Add this invoker's with warmup weight to totalWeight.
                 // If every invoker has the same weight?
                 if (sameWeight && i > 0
-                        && weight != firstWeight) {
+                        && afterWarmup != firstWeight) {
                     sameWeight = false;
                 }
             }
@@ -70,7 +68,7 @@ protected <T> Invoker<T> doSelect(List<Invoker<T>> invokers, URL url, Invocation
         }
         if (!sameWeight && totalWeight > 0) {
             // If (not every invoker has the same weight & at least one invoker's weight>0), select randomly based on totalWeight.
-            int offsetWeight = ThreadLocalRandom.current().nextInt(totalWeight);
+            int offsetWeight = ThreadLocalRandom.current().nextInt(totalWeight) + 1;
             // Return a invoker based on the random value.
             for (int i = 0; i < leastCount; i++) {
                 int leastIndex = leastIndexes[i];

dubbo-cluster/src/test/java/org/apache/dubbo/rpc/cluster/loadbalance/LeastActiveBalanceTest.java

@@ -16,13 +16,17 @@
  */
 package org.apache.dubbo.rpc.cluster.loadbalance;
 
+import org.apache.dubbo.common.Constants;
+import org.apache.dubbo.common.URL;
+import org.apache.dubbo.rpc.Invocation;
 import org.apache.dubbo.rpc.Invoker;
-
 import org.junit.Assert;
 import org.junit.Ignore;
 import org.junit.Test;
 
+import java.util.List;
 import java.util.Map;
+import java.util.Random;
 import java.util.concurrent.atomic.AtomicLong;
 
 public class LeastActiveBalanceTest extends LoadBalanceBaseTest {
@@ -39,4 +43,94 @@ public void testLeastActiveLoadBalance_select() {
         }
     }
 
+    @Test
+    public void testSelectByWeight() {
+        int sumInvoker1 = 0;
+        int sumInvoker2 = 0;
+        int loop = 100000;
+
+        MyLeastActiveLoadBalance lb = new MyLeastActiveLoadBalance();
+        for (int i = 0; i < 100000; i++) {
+            Invoker selected = lb.select(weightInvokers, null, null);
+
+            if (selected.getUrl().getProtocol().equals("test1")) {
+                sumInvoker1++;
+            }
+
+            if (selected.getUrl().getProtocol().equals("test2")) {
+                sumInvoker2++;
+            }
+            // never select invoker3 because it's active is more than invoker1 and invoker2
+            Assert.assertTrue("select is not the least active one", !selected.getUrl().getProtocol().equals("test3"));
+        }
+
+        // the sumInvoker1 : sumInvoker2 approximately equal to 1: 9
+        System.out.println(sumInvoker1);
+        System.out.println(sumInvoker2);
+
+        Assert.assertEquals("select failed!", sumInvoker1 + sumInvoker2, loop);
+    }
+
+    class MyLeastActiveLoadBalance extends AbstractLoadBalance {
+
+        private final Random random = new Random();
+
+        @Override
+        protected <T> Invoker<T> doSelect(List<Invoker<T>> invokers, URL url, Invocation invocation) {
+            int length = invokers.size(); // Number of invokers
+            int leastActive = -1; // The least active value of all invokers
+            int leastCount = 0; // The number of invokers having the same least active value (leastActive)
+            int[] leastIndexs = new int[length]; // The index of invokers having the same least active value (leastActive)
+            int totalWeightAfterWarmUp = 0; // The sum of after warmup weights
+            int firstWeightAfterWarmUp = 0; // Initial value, used for comparision
+            boolean sameWeight = true; // Every invoker has the same weight value?
+            for (int i = 0; i < length; i++) {
+                Invoker<T> invoker = invokers.get(i);
+
+                // Active number
+                int active = invoker.getUrl().getParameter("active", Constants.DEFAULT_WEIGHT);
+
+                int afterWarmup = invoker.getUrl().getParameter(Constants.WEIGHT_KEY, Constants.DEFAULT_WEIGHT);
+
+                if (leastActive == -1 || active < leastActive) { // Restart, when find a invoker having smaller least active value.
+                    leastActive = active; // Record the current least active value
+                    leastCount = 1; // Reset leastCount, count again based on current leastCount
+                    leastIndexs[0] = i; // Reset
+                    totalWeightAfterWarmUp = afterWarmup; // Reset
+                    firstWeightAfterWarmUp = afterWarmup; // Record the weight the first invoker
+                    sameWeight = true; // Reset, every invoker has the same weight value?
+                } else if (active == leastActive) { // If current invoker's active value equals with leaseActive, then accumulating.
+                    leastIndexs[leastCount++] = i; // Record index number of this invoker
+                    totalWeightAfterWarmUp += afterWarmup; // Add this invoker's after warmup weight to totalWeightAfterWarmUp.
+                    // If every invoker has the same weight?
+                    if (sameWeight && i > 0
+                            && afterWarmup != firstWeightAfterWarmUp) {
+                        sameWeight = false;
+                    }
+                }
+            }
+            // assert(leastCount > 0)
+            if (leastCount == 1) {
+                // If we got exactly one invoker having the least active value, return this invoker directly.
+                return invokers.get(leastIndexs[0]);
+            }
+            if (!sameWeight && totalWeightAfterWarmUp > 0) {
+                // If (not every invoker has the same weight & at least one invoker's weight>0), select randomly based on totalWeightAfterWarmUp.
+                int offsetWeight = random.nextInt(totalWeightAfterWarmUp) + 1;
+                // Return a invoker based on the random value.
+                for (int i = 0; i < leastCount; i++) {
+                    int leastIndex = leastIndexs[i];
+
+                    offsetWeight -= invokers.get(leastIndex).getUrl().getParameter(Constants.WEIGHT_KEY, Constants.DEFAULT_WEIGHT);
+
+                    if (offsetWeight <= 0)
+                        return invokers.get(leastIndex);
+                }
+                // assert that at most loop 'leastCount' counts
+                Assert.assertTrue("leastCount is still > 0", leastCount < 0);
+            }
+            // If all invokers have the same weight value or totalWeightAfterWarmUp=0, return evenly.
+            return invokers.get(leastIndexs[random.nextInt(leastCount)]);
+        }
+    }
 }

dubbo-cluster/src/test/java/org/apache/dubbo/rpc/cluster/loadbalance/LoadBalanceBaseTest.java

@@ -145,4 +145,40 @@ private static int calculateDefaultWarmupWeight(int uptime) {
         return AbstractLoadBalance.calculateWarmupWeight(uptime, Constants.DEFAULT_WARMUP, Constants.DEFAULT_WEIGHT);
     }
 
+    /*------------------------------------test invokers for weight---------------------------------------*/
+
+    protected List<Invoker<LoadBalanceBaseTest>> weightInvokers = new ArrayList<Invoker<LoadBalanceBaseTest>>();
+    protected Invoker<LoadBalanceBaseTest> weightInvoker1;
+    protected Invoker<LoadBalanceBaseTest> weightInvoker2;
+    protected Invoker<LoadBalanceBaseTest> weightInvoker3;
+
+    @Before
+    public void before() throws Exception {
+        weightInvoker1 = mock(Invoker.class);
+        weightInvoker2 = mock(Invoker.class);
+        weightInvoker3 = mock(Invoker.class);
+
+        URL url1 = URL.valueOf("test1://0:1/DemoService");
+        url1 = url1.addParameter(Constants.WEIGHT_KEY, 1);
+        url1 = url1.addParameter("active", 0);
+        URL url2 = URL.valueOf("test2://0:9/DemoService");
+        url2 = url2.addParameter(Constants.WEIGHT_KEY, 9);
+        url2 = url2.addParameter("active", 0);
+        URL url3 = URL.valueOf("test3://1:6/DemoService");
+        url3 = url3.addParameter(Constants.WEIGHT_KEY, 6);
+        url3 = url3.addParameter("active", 1);
+
+        given(weightInvoker1.isAvailable()).willReturn(true);
+        given(weightInvoker1.getUrl()).willReturn(url1);
+
+        given(weightInvoker2.isAvailable()).willReturn(true);
+        given(weightInvoker2.getUrl()).willReturn(url2);
+
+        given(weightInvoker3.isAvailable()).willReturn(true);
+        given(weightInvoker3.getUrl()).willReturn(url3);
+
+        weightInvokers.add(weightInvoker1);
+        weightInvokers.add(weightInvoker2);
+        weightInvokers.add(weightInvoker3);
+    }
 }

dubbo-cluster/src/test/java/org/apache/dubbo/rpc/cluster/loadbalance/RandomLoadBalanceTest.java

@@ -16,13 +16,18 @@
  */
 package org.apache.dubbo.rpc.cluster.loadbalance;
 
+import org.apache.dubbo.common.Constants;
+import org.apache.dubbo.common.URL;
+import org.apache.dubbo.rpc.Invocation;
 import org.apache.dubbo.rpc.Invoker;
 import org.apache.dubbo.rpc.RpcStatus;
 
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.List;
 import java.util.Map;
+import java.util.Random;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
@@ -54,4 +59,73 @@ public void testRandomLoadBalanceSelect() {
         Assert.assertEquals(0, counter.get(invoker5).intValue());
     }
 
+    @Test
+    public void testSelectByWeight() {
+        int sumInvoker1 = 0;
+        int sumInvoker2 = 0;
+        int sumInvoker3 = 0;
+        int loop = 100000;
+
+        MyRandomLoadBalance lb = new MyRandomLoadBalance();
+        for (int i = 0; i < loop; i++) {
+            Invoker selected = lb.select(weightInvokers, null, null);
+
+            if (selected.getUrl().getProtocol().equals("test1")) {
+                sumInvoker1++;
+            }
+
+            if (selected.getUrl().getProtocol().equals("test2")) {
+                sumInvoker2++;
+            }
+
+            if (selected.getUrl().getProtocol().equals("test3")) {
+                sumInvoker3++;
+            }
+        }
+
+        // 1 : 9 : 6
+        System.out.println(sumInvoker1);
+        System.out.println(sumInvoker2);
+        System.out.println(sumInvoker3);
+        Assert.assertEquals("select failed!", sumInvoker1 + sumInvoker2 + sumInvoker3, loop);
+    }
+
+    class MyRandomLoadBalance extends AbstractLoadBalance {
+
+        public static final String NAME = "random";
+
+        private final Random random = new Random();
+
+        @Override
+        protected <T> Invoker<T> doSelect(List<Invoker<T>> invokers, URL url, Invocation invocation) {
+            int length = invokers.size(); // Number of invokers
+            int totalWeight = 0; // The sum of weights
+            boolean sameWeight = true; // Every invoker has the same weight?
+            for (int i = 0; i < length; i++) {
+
+                // mock weight
+                int weight = invokers.get(i).getUrl().getParameter(Constants.WEIGHT_KEY, Constants.DEFAULT_WEIGHT);
+
+                totalWeight += weight; // Sum
+                if (sameWeight && i > 0
+                        && weight != invokers.get(i - 1).getUrl().getParameter(Constants.WEIGHT_KEY, Constants.DEFAULT_WEIGHT)) {
+                    sameWeight = false;
+                }
+            }
+            if (totalWeight > 0 && !sameWeight) {
+                // If (not every invoker has the same weight & at least one invoker's weight>0), select randomly based on totalWeight.
+                int offset = random.nextInt(totalWeight);
+                // Return a invoker based on the random value.
+                for (int i = 0; i < length; i++) {
+                    offset -= invokers.get(i).getUrl().getParameter(Constants.WEIGHT_KEY, Constants.DEFAULT_WEIGHT);
+                    if (offset < 0) {
+                        return invokers.get(i);
+                    }
+                }
+            }
+            // If all invokers have the same weight value or totalWeight=0, return evenly.
+            return invokers.get(random.nextInt(length));
+        }
+    }
+
 }