fixes #854 RoundRobinLoadBalance is unpredictable if there are more t…

…han two downstream APIs (#855)
networknt · Dec 7, 2020 · fcded16 · fcded16
1 parent db85287
commit fcded16
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 25 deletions.
diff --git a/balance/src/main/java/com/networknt/balance/ConsistentHashLoadBalance.java b/balance/src/main/java/com/networknt/balance/ConsistentHashLoadBalance.java
@@ -51,7 +51,7 @@ public ConsistentHashLoadBalance() {
 
 
     @Override
-    public URL select(List<URL> urls, String requestKey) {
+    public URL select(List<URL> urls, String serviceId, String tag, String requestKey) {
         URL url = null;
         if (urls.size() > 1) {
             url = doSelect(urls, requestKey);

diff --git a/balance/src/main/java/com/networknt/balance/LoadBalance.java b/balance/src/main/java/com/networknt/balance/LoadBalance.java
@@ -27,10 +27,12 @@ public interface LoadBalance {
      * Select one url from a list of url with requestKey as optional.
      *
      * @param urls List
+     * @param serviceId String
+     * @param tag String
      * @param requestKey String
      * @return URL
      */
-    URL select(List<URL> urls, String requestKey);
+    URL select(List<URL> urls, String serviceId, String tag, String requestKey);
 
     /**
      * return positive int value of originValue

diff --git a/balance/src/main/java/com/networknt/balance/LocalFirstLoadBalance.java b/balance/src/main/java/com/networknt/balance/LocalFirstLoadBalance.java
@@ -65,23 +65,26 @@ public LocalFirstLoadBalance() {
      * to use the remote interface for service to service communication.
      *
      * @param urls List
+     * @param serviceId String
+     * @param tag String
      * @param requestKey String
      * @return URL
      */
     @Override
-    public URL select(List<URL> urls, String requestKey) {
+    public URL select(List<URL> urls, String serviceId, String tag, String requestKey) {
+        String key = tag == null ? serviceId : serviceId + "|" + tag;
     	// search for a URL in the same ip first
         List<URL> localUrls = searchLocalUrls(urls, ip);
         if(localUrls.size() > 0) {
              if(localUrls.size() == 1) {
                  return localUrls.get(0);
              } else {
                 // round robin within localUrls
-                 return doSelect(localUrls);
+                 return doSelect(localUrls, key);
              }
         } else {
             // round robin within urls
-            return doSelect(urls);
+            return doSelect(urls, key);
         }
     }
 

diff --git a/balance/src/main/java/com/networknt/balance/RoundRobinLoadBalance.java b/balance/src/main/java/com/networknt/balance/RoundRobinLoadBalance.java
@@ -21,6 +21,8 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -39,35 +41,38 @@
  */
 public class RoundRobinLoadBalance implements LoadBalance {
     static Logger logger = LoggerFactory.getLogger(RoundRobinLoadBalance.class);
+    // cache the idx for each service so that the index is per service for the round robin.
+    Map<String, AtomicInteger> serviceIdx = new ConcurrentHashMap<>();
 
     public RoundRobinLoadBalance() {
         if(logger.isInfoEnabled()) logger.info("A RoundRobinLoadBalance instance is started");
     }
 
-    private AtomicInteger idx = new AtomicInteger((int)(Math.random()*10));
-
     /**
      * Round robin requestKey is not used as it should be null, the url will
      * be selected from the list base on an instance idx so every url has the
      * same priority.
      *
      * @param urls List
+     * @param serviceId String
+     * @param tag String
      * @param requestKey String
      * @return Url
      */
     @Override
-    public URL select(List<URL> urls, String requestKey) {
+    public URL select(List<URL> urls, String serviceId, String tag, String requestKey) {
         URL url = null;
         if (urls.size() > 1) {
-            url = doSelect(urls);
+            String key = tag == null ? serviceId : serviceId + "|" + tag;
+            url = doSelect(urls, key);
         } else if (urls.size() == 1) {
             url = urls.get(0);
         }
         return url;
     }
 
-    protected URL doSelect(List<URL> urls) {
-        int index = getNextPositive();
+    protected URL doSelect(List<URL> urls, String key) {
+        int index = getNextPositive(key);
         for (int i = 0; i < urls.size(); i++) {
             URL url = urls.get((i + index) % urls.size());
             if (url != null) {
@@ -78,9 +83,12 @@ protected URL doSelect(List<URL> urls) {
     }
 
     // get positive int
-    private int getNextPositive() {
+    private int getNextPositive(String key) {
+        AtomicInteger idx = serviceIdx.get(key);
+        if(idx == null) {
+            idx = new AtomicInteger((int)(Math.random()*10));
+            serviceIdx.put(key, idx);
+        }
         return getPositive(idx.incrementAndGet());
     }
-
-
 }
diff --git a/balance/src/test/java/com/networknt/balance/LocalFirstLoadBalanceTest.java b/balance/src/test/java/com/networknt/balance/LocalFirstLoadBalanceTest.java
@@ -46,7 +46,7 @@ public void testSelect() throws Exception {
         urls.add(new URLImpl("http", "127.0.0.11", 8082, "v1", new HashMap<String, String>()));
         urls.add(new URLImpl("http", "127.0.0.12", 8083, "v1", new HashMap<String, String>()));
         urls.add(new URLImpl("http", "127.0.0.115", 8084, "v1", new HashMap<String, String>()));
-        URL url = loadBalance.select(urls, null);
+        URL url = loadBalance.select(urls, "serviceId", "tag", null);
         Assert.assertEquals(url, URLImpl.valueOf("http://127.0.0.1:8081/v1"));
     }
 
@@ -59,14 +59,14 @@ public void testSelectFirstThenRoundRobin() throws Exception{
         urls.add(new URLImpl("http", "127.0.0.10", 8084, "v1", new HashMap<String, String>()));
 
         // no local host URL available, go round-robin
-        URL url = loadBalance.select(urls, null);
+        URL url = loadBalance.select(urls, "serviceId", "tag", null);
         Assert.assertTrue(urls.contains(url));
     }
 
     @Test
     public void testSelectWithEmptyList() throws Exception {
         List<URL> urls = new ArrayList<>();
-        URL url = loadBalance.select(urls, null);
+        URL url = loadBalance.select(urls, "serviceId", "tag", null);
         Assert.assertNull(url);
     }
 }
diff --git a/balance/src/test/java/com/networknt/balance/RoundRobinLoadBalanceTest.java b/balance/src/test/java/com/networknt/balance/RoundRobinLoadBalanceTest.java
@@ -42,25 +42,25 @@ public void testSelect() throws Exception {
         urls.add(new URLImpl("http", "127.0.0.1", 8083, "v1", new HashMap<String, String>()));
         urls.add(new URLImpl("http", "127.0.0.1", 8084, "v1", new HashMap<String, String>()));
         while(true) {
-            URL url = loadBalance.select(urls, null);
+            URL url = loadBalance.select(urls, "serviceId", "tag", null);
             if(url.getPort() == 8081) break;
         }
-        URL url = loadBalance.select(urls, null);
+        URL url = loadBalance.select(urls, "serviceId", "tag", null);
         Assert.assertEquals(url, URLImpl.valueOf("http://127.0.0.1:8082/v1"));
-        url = loadBalance.select(urls, null);
+        url = loadBalance.select(urls, "serviceId", "tag", null);
         Assert.assertEquals(url, URLImpl.valueOf("http://127.0.0.1:8083/v1"));
-        url = loadBalance.select(urls, null);
+        url = loadBalance.select(urls, "serviceId", "tag", null);
         Assert.assertEquals(url, URLImpl.valueOf("http://127.0.0.1:8084/v1"));
-        url = loadBalance.select(urls, null);
+        url = loadBalance.select(urls, "serviceId", "tag", null);
         Assert.assertEquals(url, URLImpl.valueOf("http://127.0.0.1:8081/v1"));
-        url = loadBalance.select(urls, null);
+        url = loadBalance.select(urls, "serviceId", "tag", null);
         Assert.assertEquals(url, URLImpl.valueOf("http://127.0.0.1:8082/v1"));
 
     }
     @Test
     public void testSelectWithEmptyList() throws Exception {
         List<URL> urls = new ArrayList<>();
-        URL url = loadBalance.select(urls, null);
+        URL url = loadBalance.select(urls, "serviceId", "tag", null);
         Assert.assertNull(url);
     }
 

diff --git a/cluster/src/main/java/com/networknt/cluster/LightCluster.java b/cluster/src/main/java/com/networknt/cluster/LightCluster.java
@@ -57,7 +57,7 @@ public LightCluster() {
      */
     @Override
     public String serviceToUrl(String protocol, String serviceId, String tag, String requestKey) {
-        URL url = loadBalance.select(discovery(protocol, serviceId, tag), requestKey);
+        URL url = loadBalance.select(discovery(protocol, serviceId, tag), serviceId, tag, requestKey);
         if (url != null) {
             logger.debug("Final url after load balance = {}.", url);
             // construct a url in string