Refactor/decouple load balance metrics collection from divide plugin

shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java

@@ -1068,6 +1068,11 @@ public interface Constants {
      */
     String GB = "GB";
 
+    /**
+     * The request begin time.
+     */
+    String REQUEST_BEGIN_TIME = "requestBeginTime";
+
     /**
      * String q.
      */

shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/factory/LoadBalancerFactory.java

@@ -44,4 +44,9 @@ public static Upstream selector(final List<Upstream> upstreamList, final String
         LoadBalancer loadBalance = ExtensionLoader.getExtensionLoader(LoadBalancer.class).getJoin(algorithm);
         return loadBalance.select(upstreamList, data);
     }
+
+    public static LoadBalancer getInstance(final String algorithm) {
+        LoadBalancer loadBalance = ExtensionLoader.getExtensionLoader(LoadBalancer.class).getJoin(algorithm);
+        return loadBalance;
+    }
 }

shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/LoadBalancer.java

@@ -33,8 +33,15 @@ public interface LoadBalancer {
      * this is select one for upstream list.
      *
      * @param upstreamList upstream list
-     * @param data data
+     * @param data         data
      * @return upstream
      */
     Upstream select(List<Upstream> upstreamList, LoadBalanceData data);
+
+    default void onSuccess(Upstream upstream, LoadBalanceData data) {
+    }
+
+    default void onError(Upstream upstream, LoadBalanceData data) {
+    }
+
 }

shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/P2cLoadBalancer.java

@@ -46,7 +46,40 @@ public class P2cLoadBalancer extends AbstractLoadBalancer {
     private static final int PICK_TIMES = 3;
 
     private final Random random = new Random();
-
+
+    @Override
+    public void onSuccess(final Upstream upstream, final LoadBalanceData data) {
+        responseTrigger(upstream);
+    }
+
+    @Override
+    public void onError(final Upstream upstream, final LoadBalanceData data) {
+        responseTrigger(upstream);
+    }
+
+    private void responseTrigger(final Upstream upstream) {
+        long now = System.currentTimeMillis();
+        upstream.getInflight().decrementAndGet();
+        long stamp = upstream.getResponseStamp();
+        upstream.setResponseStamp(now);
+        long td = now - stamp;
+        if (td < 0) {
+            td = 0;
+        }
+        double w = Math.exp((double) -td / (double) 600);
+
+        long lag = now - upstream.getLastPicked();
+        if (lag < 0) {
+            lag = 0;
+        }
+        long oldLag = upstream.getLag();
+        if (oldLag == 0) {
+            w = 0;
+        }
+        lag = (int) ((double) oldLag * w + (double) lag * (1.0 - w));
+        upstream.setLag(lag);
+    }
+
     @Override
     protected Upstream doSelect(final List<Upstream> upstreamList, final LoadBalanceData data) {
         long start = System.currentTimeMillis();

shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/ShortestResponseLoadBalancer.java

@@ -17,6 +17,7 @@
 
 package org.apache.shenyu.loadbalancer.spi;
 
+import org.apache.shenyu.common.constant.Constants;
 import org.apache.shenyu.loadbalancer.entity.LoadBalanceData;
 import org.apache.shenyu.loadbalancer.entity.Upstream;
 import org.apache.shenyu.spi.Join;
@@ -81,4 +82,11 @@ protected Upstream doSelect(final List<Upstream> upstreamList, final LoadBalance
         }
         return upstreamList.get(shortestIndexes[ThreadLocalRandom.current().nextInt(shortestCount)]);
     }
+
+    @Override
+    public void onSuccess(final Upstream upstream, final LoadBalanceData data) {
+        long beginTime = (long) data.getAttributes().get(Constants.REQUEST_BEGIN_TIME);
+        upstream.getSucceededElapsed().addAndGet(System.currentTimeMillis() - beginTime);
+        upstream.getSucceeded().incrementAndGet();
+    }
 }

shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/factory/LoadBalancerFactoryTest.java

@@ -22,6 +22,8 @@
 import org.apache.shenyu.loadbalancer.entity.Upstream;
 import org.junit.jupiter.api.Test;
 
+import org.apache.shenyu.loadbalancer.spi.LoadBalancer;
+
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -30,6 +32,8 @@
 import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 
 /**
  * The type loadBalance utils test.
@@ -92,4 +96,12 @@ public void loadBalanceUtilsReversedWeightTest() {
         });
         assertEquals(12, countMap.get("upstream-10").intValue());
     }
+
+    @Test
+    public void testGetInstance() {
+        LoadBalancer lb = LoadBalancerFactory.getInstance(LoadBalanceEnum.RANDOM.getName());
+        assertNotNull(lb);
+        LoadBalancer lb2 = LoadBalancerFactory.getInstance(LoadBalanceEnum.RANDOM.getName());
+        assertSame(lb, lb2);
+    }
 }

shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/spi/P2cLoadBalancerTest.java

@@ -63,4 +63,24 @@ public void testResponseTimeBalancerSameInflight() {
         Upstream upstream1 = p2cLoadBalancer.doSelect(upstreamList, new LoadBalanceData());
         Assertions.assertTrue(upstream.getUrl().equals("baidu.com") && upstream1.getUrl().equals("pro.jd.com"));
     }
+
+    @Test
+    public void testOnSuccess() {
+        buildUpstreamList();
+        final P2cLoadBalancer p2cLoadBalancer = new P2cLoadBalancer();
+        Upstream upstream = upstreamList.get(0);
+        long beforeInflight = upstream.getInflight().get();
+        p2cLoadBalancer.onSuccess(upstream, new LoadBalanceData());
+        Assertions.assertEquals(beforeInflight - 1, upstream.getInflight().get());
+    }
+
+    @Test
+    public void testOnError() {
+        buildUpstreamList();
+        final P2cLoadBalancer p2cLoadBalancer = new P2cLoadBalancer();
+        Upstream upstream = upstreamList.get(0);
+        long beforeInflight = upstream.getInflight().get();
+        p2cLoadBalancer.onError(upstream, new LoadBalanceData());
+        Assertions.assertEquals(beforeInflight - 1, upstream.getInflight().get());
+    }
 }

shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/spi/ShortestResponseLoadBalancerTest.java

@@ -17,12 +17,15 @@
 
 package org.apache.shenyu.loadbalancer.spi;
 
+import org.apache.shenyu.common.constant.Constants;
 import org.apache.shenyu.loadbalancer.entity.LoadBalanceData;
 import org.apache.shenyu.loadbalancer.entity.Upstream;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -84,4 +87,17 @@ public void testSelectByResponse() {
         Assertions.assertEquals(loop, select1 + select2);
 
     }
+
+    @Test
+    public void testOnSuccess() {
+        Upstream upstream = Upstream.builder().url("upstream-1").build();
+        LoadBalanceData data = new LoadBalanceData();
+        Map<String, Object> attributes = new HashMap<>();
+        attributes.put(Constants.REQUEST_BEGIN_TIME, System.currentTimeMillis() - 100);
+        data.setAttributes(attributes);
+        ShortestResponseLoadBalancer lb = new ShortestResponseLoadBalancer();
+        lb.onSuccess(upstream, data);
+        Assertions.assertTrue(upstream.getSucceededElapsed().get() >= 100);
+        Assertions.assertEquals(1, upstream.getSucceeded().get());
+    }
 }

shenyu-plugin/shenyu-plugin-base/src/main/java/org/apache/shenyu/plugin/base/utils/LoadbalancerUtils.java

@@ -65,8 +65,8 @@ public static Upstream getForExchange(final List<Upstream> upstreamList, final S
     public static Upstream getForNoExchange(final List<Upstream> upstreamList, final String algorithm) {
         return LoadBalancerFactory.selector(upstreamList, algorithm, new LoadBalanceData());
     }
-    
-    private static LoadBalanceData buildLoadBalanceData(final ServerWebExchange exchange) {
+
+    public static LoadBalanceData buildLoadBalanceData(final ServerWebExchange exchange) {
         ServerHttpRequest request = exchange.getRequest();
         String ip = Objects.requireNonNull(request.getRemoteAddress()).getAddress().getHostAddress();
         String httpMethod = request.getMethod().name();

shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-divide/src/main/java/org/apache/shenyu/plugin/divide/DividePlugin.java

@@ -28,7 +28,10 @@
 import org.apache.shenyu.common.enums.RetryEnum;
 import org.apache.shenyu.common.enums.RpcTypeEnum;
 import org.apache.shenyu.loadbalancer.cache.UpstreamCacheManager;
+import org.apache.shenyu.loadbalancer.entity.LoadBalanceData;
 import org.apache.shenyu.loadbalancer.entity.Upstream;
+import org.apache.shenyu.loadbalancer.factory.LoadBalancerFactory;
+import org.apache.shenyu.loadbalancer.spi.LoadBalancer;
 import org.apache.shenyu.plugin.api.ShenyuPluginChain;
 import org.apache.shenyu.plugin.api.context.ShenyuContext;
 import org.apache.shenyu.plugin.api.result.ShenyuResultEnum;
@@ -56,12 +59,6 @@ public class DividePlugin extends AbstractShenyuPlugin {
 
     private static final Logger LOG = LoggerFactory.getLogger(DividePlugin.class);
 
-    private static final String P2C = "p2c";
-
-    private static final String SHORTEST_RESPONSE = "shortestResponse";
-
-    private Long beginTime;
-
     @Override
     protected String getRawPath(final ServerWebExchange exchange) {
         return RequestUrlUtils.getRewrittenRawPath(exchange);
@@ -130,15 +127,11 @@ protected Mono<Void> doExecute(final ServerWebExchange exchange, final ShenyuPlu
         exchange.getAttributes().put(Constants.RETRY_STRATEGY, StringUtils.defaultIfEmpty(ruleHandle.getRetryStrategy(), RetryEnum.CURRENT.getName()));
         exchange.getAttributes().put(Constants.LOAD_BALANCE, StringUtils.defaultIfEmpty(ruleHandle.getLoadBalance(), LoadBalanceEnum.RANDOM.getName()));
         exchange.getAttributes().put(Constants.DIVIDE_SELECTOR_ID, selector.getId());
-        if (ruleHandle.getLoadBalance().equals(P2C)) {
-            return chain.execute(exchange).doOnSuccess(e -> responseTrigger(upstream
-            )).doOnError(throwable -> responseTrigger(upstream));
-        } else if (ruleHandle.getLoadBalance().equals(SHORTEST_RESPONSE)) {
-            beginTime = System.currentTimeMillis();
-            return chain.execute(exchange).doOnSuccess(e -> successResponseTrigger(upstream
-            ));
-        }
-        return chain.execute(exchange);
+        exchange.getAttributes().put(Constants.REQUEST_BEGIN_TIME, System.currentTimeMillis());
+        LoadBalancer loadBalancer = LoadBalancerFactory.getInstance(ruleHandle.getLoadBalance());
+        LoadBalanceData loadBalanceData = LoadbalancerUtils.buildLoadBalanceData(exchange);
+        return chain.execute(exchange).doOnSuccess(e -> loadBalancer.onSuccess(upstream, loadBalanceData))
+                .doOnError(t -> loadBalancer.onError(upstream, loadBalanceData));
     }
 
     @Override
@@ -170,32 +163,4 @@ private DivideRuleHandle buildRuleHandle(final RuleData rule) {
         return DividePluginDataHandler.CACHED_HANDLE.get().obtainHandle(CacheKeyUtils.INST.getKey(rule));
     }
 
-    private void responseTrigger(final Upstream upstream) {
-        long now = System.currentTimeMillis();
-        upstream.getInflight().decrementAndGet();
-        upstream.setResponseStamp(now);
-        long stamp = upstream.getResponseStamp();
-        long td = now - stamp;
-        if (td < 0) {
-            td = 0;
-        }
-        double w = Math.exp((double) -td / (double) 600);
-
-        long lag = now - upstream.getLastPicked();
-        if (lag < 0) {
-            lag = 0;
-        }
-        long oldLag = upstream.getLag();
-        if (oldLag == 0) {
-            w = 0;
-        }
-        lag = (int) ((double) oldLag * w + (double) lag * (1.0 - w));
-        upstream.setLag(lag);
-    }
-
-    private void successResponseTrigger(final Upstream upstream) {
-        upstream.getSucceededElapsed().addAndGet(System.currentTimeMillis() - beginTime);
-        upstream.getSucceeded().incrementAndGet();
-    }
-
 }

shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-divide/src/test/java/org/apache/shenyu/plugin/divide/DividePluginTest.java

@@ -28,8 +28,8 @@
 import org.apache.shenyu.common.utils.GsonUtils;
 import org.apache.shenyu.common.utils.UpstreamCheckUtils;
 import org.apache.shenyu.loadbalancer.cache.UpstreamCacheManager;
-import org.apache.shenyu.loadbalancer.entity.Upstream;
 import org.apache.shenyu.loadbalancer.factory.LoadBalancerFactory;
+import org.apache.shenyu.loadbalancer.spi.LoadBalancer;
 import org.apache.shenyu.plugin.api.ShenyuPluginChain;
 import org.apache.shenyu.plugin.api.context.ShenyuContext;
 import org.apache.shenyu.plugin.api.result.DefaultShenyuResult;
@@ -53,16 +53,12 @@
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
-import java.lang.reflect.Field;
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
 import java.net.InetSocketAddress;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -160,9 +156,13 @@ public void doExecuteTest() {
         dividePlugin.doExecute(exchange, chain, selectorData, ruleData);
         // hit `Objects.isNull(upstream)`
         MockedStatic<LoadBalancerFactory> loadBalancerFactoryMockedStatic = mockStatic(LoadBalancerFactory.class);
+        LoadBalancer mockLoadBalancer = mock(LoadBalancer.class);
         loadBalancerFactoryMockedStatic.when(() -> LoadBalancerFactory.selector(any(), any(), any()))
                 .thenReturn(null);
+        loadBalancerFactoryMockedStatic.when(() -> LoadBalancerFactory.getInstance(any()))
+                .thenReturn(mockLoadBalancer);
         dividePlugin.doExecute(exchange, chain, selectorData, ruleData);
+        loadBalancerFactoryMockedStatic.close();
         // hit `Objects.requireNonNull(shenyuContext)`
         exchange.getAttributes().remove(Constants.CONTEXT);
         assertThrows(NullPointerException.class, () -> dividePlugin.doExecute(exchange, chain, selectorData, ruleData));
@@ -218,33 +218,6 @@ public void getOrderTest() {
         assertEquals(PluginEnum.DIVIDE.getCode(), dividePlugin.getOrder());
     }
 
-    @Test
-    public void responseTriggerTest() throws NoSuchMethodException, InstantiationException, IllegalAccessException, InvocationTargetException {
-        Upstream upstream = Upstream.builder()
-                .url("upstream")
-                .build();
-        assertEquals(0, upstream.getLag());
-        Method method = DividePlugin.class.getDeclaredMethod("responseTrigger", Upstream.class);
-        method.setAccessible(true);
-        method.invoke(DividePlugin.class.newInstance(), upstream);
-        assertNotEquals(0, upstream.getLag());
-    }
-
-    @Test
-    public void successResponseTriggerTest() throws Exception {
-        dividePlugin = DividePlugin.class.newInstance();
-        Field field = DividePlugin.class.getDeclaredField("beginTime");
-        field.setAccessible(true);
-        field.set(dividePlugin, 0L);
-        Method method = DividePlugin.class.getDeclaredMethod("successResponseTrigger", Upstream.class);
-        method.setAccessible(true);
-        Upstream upstream = Upstream.builder()
-                .url("upstream")
-                .build();
-        method.invoke(dividePlugin, upstream);
-        assertEquals(1, upstream.getSucceeded().get());
-    }
-
     /**
      * Init mock info.
      */