Fix load balancing (#3034)

* Fix load balancing
zowe · Aug 21, 2023 · ce2f280 · ce2f280
1 parent 52ec39b
commit ce2f280
Show file tree

Hide file tree

Showing 5 changed files with 329 additions and 80 deletions.
diff --git a/gateway-service/src/main/java/org/zowe/apiml/gateway/ribbon/ApimlRibbonRetryFactory.java b/gateway-service/src/main/java/org/zowe/apiml/gateway/ribbon/ApimlRibbonRetryFactory.java
@@ -10,57 +10,65 @@
 
 package org.zowe.apiml.gateway.ribbon;
 
-import lombok.RequiredArgsConstructor;
-import lombok.experimental.Delegate;
 import org.springframework.cloud.client.loadbalancer.LoadBalancedRetryContext;
 import org.springframework.cloud.client.loadbalancer.LoadBalancedRetryPolicy;
 import org.springframework.cloud.client.loadbalancer.ServiceInstanceChooser;
 import org.springframework.cloud.netflix.ribbon.RibbonLoadBalancedRetryFactory;
+import org.springframework.cloud.netflix.ribbon.RibbonLoadBalancedRetryPolicy;
+import org.springframework.cloud.netflix.ribbon.RibbonLoadBalancerContext;
 import org.springframework.cloud.netflix.ribbon.SpringClientFactory;
 import org.springframework.retry.RetryListener;
 
+import java.lang.reflect.InvocationTargetException;
+import java.net.ConnectException;
 import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Allows adding RetryListeners to Ribbon Retry
  */
 public class ApimlRibbonRetryFactory extends RibbonLoadBalancedRetryFactory {
 
+    private final SpringClientFactory clientFactory;
     private final AtomicReference<ServiceInstanceChooser> serviceInstanceChooser = new AtomicReference<>();
 
     public ApimlRibbonRetryFactory(SpringClientFactory clientFactory) {
         super(clientFactory);
+        this.clientFactory = clientFactory;
     }
 
     @Override
     public LoadBalancedRetryPolicy createRetryPolicy(String service, ServiceInstanceChooser serviceInstanceChooser) {
         this.serviceInstanceChooser.set(serviceInstanceChooser);
-        return new LoadBalancedRetryPolicyFix(super.createRetryPolicy(service, serviceInstanceChooser));
-    }
+        RibbonLoadBalancerContext lbContext = this.clientFactory.getLoadBalancerContext(service);
+        return new RibbonLoadBalancedRetryPolicy(service, lbContext, serviceInstanceChooser, clientFactory.getClientConfig(service)) {
 
-    @Override
-    public RetryListener[] createRetryListeners(String service) {
-        return new RetryListener[] {
-            new InitializingRetryListener(this.serviceInstanceChooser.get()),
-            new AbortingRetryListener()
-        };
-    }
+            @Override
+            public boolean canRetry(LoadBalancedRetryContext context) {
+                return super.canRetry(context) || isConnectionRefused(context.getLastThrowable());
+            }
 
-    @RequiredArgsConstructor
-    private static class LoadBalancedRetryPolicyFix implements LoadBalancedRetryPolicy {
+            @Override
+            public boolean canRetryNextServer(LoadBalancedRetryContext context) {
+                return super.canRetryNextServer(context) || context.getRetryCount() == 0;
+            }
 
-        @Delegate(excludes = CanRetryNextServer.class)
-        private final LoadBalancedRetryPolicy original;
+            private boolean isConnectionRefused(Throwable t) {
+                if (t instanceof InvocationTargetException) {
+                    return isConnectionRefused(((InvocationTargetException) t).getTargetException());
+                }
 
-        @Override
-        public boolean canRetryNextServer(LoadBalancedRetryContext context) {
-            return original.canRetryNextServer(context) || context.getRetryCount() == 0;
-        }
+                return t instanceof ConnectException;
+            }
 
-        interface CanRetryNextServer {
-            boolean canRetryNextServer(LoadBalancedRetryContext context);
-        }
+        };
+    }
 
+    @Override
+    public RetryListener[] createRetryListeners(String service) {
+        return new RetryListener[]{
+                new InitializingRetryListener(this.serviceInstanceChooser.get()),
+                new AbortingRetryListener()
+        };
     }
 
 }
diff --git a/...ice/src/main/java/org/zowe/apiml/gateway/ribbon/loadbalancer/LoadBalancerRuleAdapter.java b/...ice/src/main/java/org/zowe/apiml/gateway/ribbon/loadbalancer/LoadBalancerRuleAdapter.java
@@ -18,10 +18,13 @@
 import lombok.extern.slf4j.Slf4j;
 import org.zowe.apiml.gateway.context.ConfigurableNamedContextFactory;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import static java.util.Comparator.comparing;
+
 /**
  * This adapter holds the load balancing logic by facilitating server selection.
  * There is plenty of debug log to increase supportability
@@ -34,7 +37,7 @@ public class LoadBalancerRuleAdapter extends ClientConfigEnabledRoundRobinRule {
     private ConfigurableNamedContextFactory<?> configurableNamedContextFactory;
     private Map<String, RequestAwarePredicate> predicateMap;
 
-    // used zuul's implementation of round robin server selection
+    // used zuul's implementation of round-robin server selection
     private AvailabilityPredicate availabilityPredicate;
     private AbstractServerPredicate zuulPredicate;
 
@@ -63,7 +66,8 @@ public Server choose(Object key) {
         log.debug("Choosing server: {}", key);
         ILoadBalancer lb = getLoadBalancer();
         LoadBalancingContext ctx = new LoadBalancingContext(instanceInfo.getAppName(), instanceInfo);
-        List<Server> allServers = lb.getAllServers();
+        List<Server> allServers = new ArrayList<>(lb.getAllServers());
+        allServers.sort(comparing(Server::isReadyToServe).reversed().thenComparing(Server::getId)); // the original list is in the random order
         log.debug("Path: {}, List of servers from LoadBalancer: {}", ctx.getPath() ,allServers);
         for (RequestAwarePredicate predicate : predicateMap.values()) {
             log.debug("Running predicate: {}, list of servers: {}", allServers, predicate);

diff --git a/...src/test/java/org/zowe/apiml/gateway/ribbon/loadbalancer/LoadBalancerRuleAdapterTest.java b/...src/test/java/org/zowe/apiml/gateway/ribbon/loadbalancer/LoadBalancerRuleAdapterTest.java
@@ -14,28 +14,29 @@
 import com.netflix.loadbalancer.ILoadBalancer;
 import com.netflix.loadbalancer.Server;
 import com.netflix.niws.loadbalancer.DiscoveryEnabledServer;
-import org.junit.jupiter.api.*;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+import org.springframework.cloud.context.named.NamedContextFactory;
 import org.zowe.apiml.gateway.context.ConfigurableNamedContextFactory;
 
 import java.util.*;
 
 import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.*;
 
 class LoadBalancerRuleAdapterTest {
 
     private static DiscoveryEnabledServer server;
     private static DiscoveryEnabledServer server1;
-    private static DiscoveryEnabledServer server2;
     private static ILoadBalancer lb;
 
     @BeforeAll
     static void setup() {
         lb = mock(ILoadBalancer.class);
-        server = mock(DiscoveryEnabledServer.class);
-        server1 = mock(DiscoveryEnabledServer.class);
-        server2 = mock(DiscoveryEnabledServer.class);
+        server = createServer("server");
+        server1 = createServer("server2");
     }
 
     @Nested
@@ -44,7 +45,9 @@ class GivenOnlyDefaultPredicate {
         void choosesRoundRobin() {
             LoadBalancerRuleAdapter underTest = new LoadBalancerRuleAdapter(mock(InstanceInfo.class), mock(ConfigurableNamedContextFactory.class), null);
             underTest.setLoadBalancer(lb);
-            when(lb.getAllServers()).thenReturn(Arrays.asList(server, server1));
+            List<Server> serverList = Arrays.asList(server, server1);
+            Collections.shuffle(serverList);
+            when(lb.getAllServers()).thenReturn(serverList);
             Server theChosenOne = underTest.choose("key");
             Server theChosenTwo = underTest.choose("key");
             Server theChosenThree = underTest.choose("key");
@@ -59,7 +62,7 @@ void choosesRoundRobin() {
     @Nested
     class GivenAdditionalPredicate {
 
-        private ConfigurableNamedContextFactory configurableNamedContextFactory;
+        private ConfigurableNamedContextFactory<NamedContextFactory.Specification> configurableNamedContextFactory;
         private RequestAwarePredicate requestAwarePredicate;
         private RequestAwarePredicate requestAwarePredicate1;
         private InstanceInfo instanceInfo;
@@ -107,10 +110,11 @@ void noServerFitsThePredicate() {
     }
 
     @Nested
-    class givenHeterogenousListOfServers {
+    class givenHeterogeneousListOfServers {
+
 
-        private Map predicateMap = new HashMap<>();
-        private ConfigurableNamedContextFactory configurableNamedContextFactory = mock(ConfigurableNamedContextFactory.class);
+        private final Map<String, Object> predicateMap = new HashMap<>();
+        private final ConfigurableNamedContextFactory<NamedContextFactory.Specification> configurableNamedContextFactory = mock(ConfigurableNamedContextFactory.class);
 
         @Test
         void shouldFailFast() {
@@ -124,4 +128,12 @@ void shouldFailFast() {
         }
     }
 
+    private static DiscoveryEnabledServer createServer(String name) {
+        InstanceInfo instanceInfo = InstanceInfo.Builder.newBuilder()
+                .setAppName(name)
+                .setHostName(name)
+                .build();
+        return new DiscoveryEnabledServer(instanceInfo, true);
+    }
+
 }