Optimize RoundRobinLoadBalance

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

@@ -21,83 +21,64 @@
 import org.apache.dubbo.rpc.Invocation;
 import org.apache.dubbo.rpc.Invoker;
 
-import java.util.LinkedHashMap;
+import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
 /**
  * Round robin load balance.
- *
  */
 public class RoundRobinLoadBalance extends AbstractLoadBalance {
 
     public static final String NAME = "roundrobin";
 
     private final ConcurrentMap<String, AtomicPositiveInteger> sequences = new ConcurrentHashMap<String, AtomicPositiveInteger>();
 
+    private final ConcurrentMap<String, AtomicPositiveInteger> indexSeqs = new ConcurrentHashMap<String, AtomicPositiveInteger>();
+
     @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;
+        final List<Invoker<T>> invokerToWeightList = new ArrayList<>();
         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;
+                invokerToWeightList.add(invokers.get(i));
             }
         }
         AtomicPositiveInteger sequence = sequences.get(key);
         if (sequence == null) {
             sequences.putIfAbsent(key, new AtomicPositiveInteger());
             sequence = sequences.get(key);
         }
-        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--;
-                    }
+            AtomicPositiveInteger indexSeq = indexSeqs.get(key);
+            if (indexSeq == null) {
+                indexSeqs.putIfAbsent(key, new AtomicPositiveInteger(-1));
+                indexSeq = indexSeqs.get(key);
+            }
+            length = invokerToWeightList.size();
+            while (true) {
+                int index = indexSeq.incrementAndGet() % length;
+                int currentWeight;
+                if (index == 0) {
+                    currentWeight = sequence.incrementAndGet() % maxWeight;
+                }else {
+                    currentWeight = sequence.get() % maxWeight;
+                }
+                if (getWeight(invokerToWeightList.get(index), invocation) > currentWeight) {
+                    return invokerToWeightList.get(index);
                 }
             }
         }
         // 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--;
-        }
+        return invokers.get(sequence.getAndIncrement() % length);
     }
-
 }