Smooth Round Robin selection

Author: jasonjoo2010

dubbo-cluster/src/main/java/com/alibaba/dubbo/rpc/cluster/loadbalance/RoundRobinLoadBalance.java

@@ -17,87 +17,141 @@
 package com.alibaba.dubbo.rpc.cluster.loadbalance;
 
 import com.alibaba.dubbo.common.URL;
-import com.alibaba.dubbo.common.utils.AtomicPositiveInteger;
 import com.alibaba.dubbo.rpc.Invocation;
 import com.alibaba.dubbo.rpc.Invoker;
 
-import java.util.LinkedHashMap;
+import java.util.Collection;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Round robin load balance.
+ * 
+ * Smoothly round robin's implementation @since 2.6.5 
+ * @author jason
  *
  */
 public class RoundRobinLoadBalance extends AbstractLoadBalance {
-
     public static final String NAME = "roundrobin";
+    
+    private static int RECYCLE_PERIOD = 60000;
+    
+    protected static class WeightedRoundRobin {
+        private int weight;
+        private AtomicLong current = new AtomicLong(0);
+        private long lastUpdate;
+        public int getWeight() {
+            return weight;
+        }
+        public void setWeight(int weight) {
+            this.weight = weight;
+            current.set(0);
+        }
+        public long increaseCurrent() {
+            return current.addAndGet(weight);
+        }
+        public void sel(int total) {
+            current.addAndGet(-1 * total);
+        }
+        public long getLastUpdate() {
+            return lastUpdate;
+        }
+        public void setLastUpdate(long lastUpdate) {
+            this.lastUpdate = lastUpdate;
+        }
+    }
 
-    private final ConcurrentMap<String, AtomicPositiveInteger> sequences = new ConcurrentHashMap<String, AtomicPositiveInteger>();
-
+    private ConcurrentMap<String, ConcurrentMap<String, WeightedRoundRobin>> methodWeightMap = new ConcurrentHashMap<String, ConcurrentMap<String, WeightedRoundRobin>>();
+    private AtomicBoolean updateLock = new AtomicBoolean();
+    
+    /**
+     * get invoker addr list cached for specified invocation
+     * <p>
+     * <b>for unit test only</b>
+     * 
+     * @param invokers
+     * @param invocation
+     * @return
+     */
+    protected <T> Collection<String> getInvokerAddrList(List<Invoker<T>> invokers, Invocation invocation) {
+        String key = invokers.get(0).getUrl().getServiceKey() + "." + invocation.getMethodName();
+        Map<String, WeightedRoundRobin> map = methodWeightMap.get(key);
+        if (map != null) {
+            return map.keySet();
+        }
+        return null;
+    }
+    
     @Override
     protected <T> Invoker<T> doSelect(List<Invoker<T>> invokers, URL url, Invocation invocation) {
         String key = invokers.get(0).getUrl().getServiceKey() + "." + invocation.getMethodName();
-        int length = invokers.size(); // Number of invokers
-        int maxWeight = 0; // The maximum weight
-        int minWeight = Integer.MAX_VALUE; // The minimum weight
-        final LinkedHashMap<Invoker<T>, IntegerWrapper> invokerToWeightMap = new LinkedHashMap<Invoker<T>, IntegerWrapper>();
-        int weightSum = 0;
-        for (int i = 0; i < length; i++) {
-            int weight = getWeight(invokers.get(i), invocation);
-            maxWeight = Math.max(maxWeight, weight); // Choose the maximum weight
-            minWeight = Math.min(minWeight, weight); // Choose the minimum weight
-            if (weight > 0) {
-                invokerToWeightMap.put(invokers.get(i), new IntegerWrapper(weight));
-                weightSum += weight;
-            }
+        ConcurrentMap<String, WeightedRoundRobin> map = methodWeightMap.get(key);
+        if (map == null) {
+            methodWeightMap.putIfAbsent(key, new ConcurrentHashMap<String, WeightedRoundRobin>());
+            map = methodWeightMap.get(key);
         }
-        AtomicPositiveInteger sequence = sequences.get(key);
-        if (sequence == null) {
-            sequences.putIfAbsent(key, new AtomicPositiveInteger());
-            sequence = sequences.get(key);
+        int totalWeight = 0;
+        long maxCurrent = Long.MIN_VALUE;
+        long now = System.currentTimeMillis();
+        Invoker<T> selectedInvoker = null;
+        WeightedRoundRobin selectedWRR = null;
+        for (Invoker<T> invoker : invokers) {
+            String identifyString = invoker.getUrl().toIdentityString();
+            WeightedRoundRobin weightedRoundRobin = map.get(identifyString);
+            int weight = getWeight(invoker, invocation);
+            if (weight < 0) {
+                weight = 0;
+            }
+            if (weightedRoundRobin == null) {
+                weightedRoundRobin = new WeightedRoundRobin();
+                weightedRoundRobin.setWeight(weight);
+                map.putIfAbsent(identifyString, weightedRoundRobin);
+                weightedRoundRobin = map.get(identifyString);
+            }
+            if (weight != weightedRoundRobin.getWeight()) {
+                //weight changed
+                weightedRoundRobin.setWeight(weight);
+            }
+            long cur = weightedRoundRobin.increaseCurrent();
+            weightedRoundRobin.setLastUpdate(now);
+            if (cur > maxCurrent) {
+                maxCurrent = cur;
+                selectedInvoker = invoker;
+                selectedWRR = weightedRoundRobin;
+            }
+            totalWeight += weight;
         }
-        int currentSequence = sequence.getAndIncrement();
-        if (maxWeight > 0 && minWeight < maxWeight) {
-            int mod = currentSequence % weightSum;
-            for (int i = 0; i < maxWeight; i++) {
-                for (Map.Entry<Invoker<T>, IntegerWrapper> each : invokerToWeightMap.entrySet()) {
-                    final Invoker<T> k = each.getKey();
-                    final IntegerWrapper v = each.getValue();
-                    if (mod == 0 && v.getValue() > 0) {
-                        return k;
-                    }
-                    if (v.getValue() > 0) {
-                        v.decrement();
-                        mod--;
+        if (!updateLock.get() && invokers.size() != map.size()) {
+            if (updateLock.compareAndSet(false, true)) {
+                try {
+                    // copy -> modify -> update reference
+                    ConcurrentMap<String, WeightedRoundRobin> newMap = new ConcurrentHashMap<String, WeightedRoundRobin>();
+                    newMap.putAll(map);
+                    Iterator<Entry<String, WeightedRoundRobin>> it = newMap.entrySet().iterator();
+                    while (it.hasNext()) {
+                        Entry<String, WeightedRoundRobin> item = it.next();
+                        if (now - item.getValue().getLastUpdate() > RECYCLE_PERIOD) {
+                            it.remove();
+                        }
                     }
+                    methodWeightMap.put(key, newMap);
+                } finally {
+                    updateLock.set(false);
                 }
             }
         }
-        // Round robin
-        return invokers.get(currentSequence % length);
-    }
-
-    private static final class IntegerWrapper {
-        private int value;
-
-        public IntegerWrapper(int value) {
-            this.value = value;
-        }
-
-        public int getValue() {
-            return value;
-        }
-
-        public void setValue(int value) {
-            this.value = value;
-        }
-
-        public void decrement() {
-            this.value--;
+        if (selectedInvoker != null) {
+            selectedWRR.sel(totalWeight);
+            return selectedInvoker;
         }
+        // should not happen here
+        return invokers.get(0);
     }
 
 }

dubbo-cluster/src/test/java/com/alibaba/dubbo/rpc/cluster/StickyTest.java

@@ -114,12 +114,12 @@ public int testSticky(String methodName, boolean check) {
 
         given(invoker1.invoke(invocation)).willReturn(result);
         given(invoker1.isAvailable()).willReturn(true);
-        given(invoker1.getUrl()).willReturn(url);
+        given(invoker1.getUrl()).willReturn(url.setPort(1));
         given(invoker1.getInterface()).willReturn(StickyTest.class);
 
         given(invoker2.invoke(invocation)).willReturn(result);
         given(invoker2.isAvailable()).willReturn(true);
-        given(invoker2.getUrl()).willReturn(url);
+        given(invoker2.getUrl()).willReturn(url.setPort(2));
         given(invoker2.getInterface()).willReturn(StickyTest.class);
 
         invocation.setMethodName(methodName);
@@ -158,4 +158,4 @@ public Invoker<T> getSelectedInvoker() {
             return selectedInvoker;
         }
     }
-}
+}