diff --git a/shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java b/shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java index 5295c8552b8c..eec095b10b8c 100644 --- a/shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java +++ b/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. */ diff --git a/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/factory/LoadBalancerFactory.java b/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/factory/LoadBalancerFactory.java index 291039da43ea..1a2370cbc681 100644 --- a/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/factory/LoadBalancerFactory.java +++ b/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/factory/LoadBalancerFactory.java @@ -44,4 +44,9 @@ public static Upstream selector(final List 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; + } } diff --git a/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/LoadBalancer.java b/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/LoadBalancer.java index ad0d0de3e2a7..de85a9103a56 100644 --- a/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/LoadBalancer.java +++ b/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 upstreamList, LoadBalanceData data); + + default void onSuccess(Upstream upstream, LoadBalanceData data) { + } + + default void onError(Upstream upstream, LoadBalanceData data) { + } + } diff --git a/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/P2cLoadBalancer.java b/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/P2cLoadBalancer.java index dd0edb7cebc3..c56215374e30 100644 --- a/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/P2cLoadBalancer.java +++ b/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 upstreamList, final LoadBalanceData data) { long start = System.currentTimeMillis(); diff --git a/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/ShortestResponseLoadBalancer.java b/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/ShortestResponseLoadBalancer.java index edf4f15aa6cb..b869a2ab343a 100644 --- a/shenyu-loadbalancer/src/main/java/org/apache/shenyu/loadbalancer/spi/ShortestResponseLoadBalancer.java +++ b/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 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(); + } } diff --git a/shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/factory/LoadBalancerFactoryTest.java b/shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/factory/LoadBalancerFactoryTest.java index e7b7f2f280f9..6f8e14e84a07 100644 --- a/shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/factory/LoadBalancerFactoryTest.java +++ b/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); + } } diff --git a/shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/spi/P2cLoadBalancerTest.java b/shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/spi/P2cLoadBalancerTest.java index 5b760949fc2d..2f31db035157 100644 --- a/shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/spi/P2cLoadBalancerTest.java +++ b/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()); + } } diff --git a/shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/spi/ShortestResponseLoadBalancerTest.java b/shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/spi/ShortestResponseLoadBalancerTest.java index b6bb359d20b7..854a0837735e 100644 --- a/shenyu-loadbalancer/src/test/java/org/apache/shenyu/loadbalancer/spi/ShortestResponseLoadBalancerTest.java +++ b/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 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()); + } } diff --git a/shenyu-plugin/shenyu-plugin-base/src/main/java/org/apache/shenyu/plugin/base/utils/LoadbalancerUtils.java b/shenyu-plugin/shenyu-plugin-base/src/main/java/org/apache/shenyu/plugin/base/utils/LoadbalancerUtils.java index 91e8e850ab6a..5ae191f087be 100644 --- a/shenyu-plugin/shenyu-plugin-base/src/main/java/org/apache/shenyu/plugin/base/utils/LoadbalancerUtils.java +++ b/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 upstreamList, final S public static Upstream getForNoExchange(final List 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(); diff --git a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-divide/src/main/java/org/apache/shenyu/plugin/divide/DividePlugin.java b/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-divide/src/main/java/org/apache/shenyu/plugin/divide/DividePlugin.java index 5505c8fdbbda..7f5227897c1f 100644 --- a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-divide/src/main/java/org/apache/shenyu/plugin/divide/DividePlugin.java +++ b/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 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(); - } - } diff --git a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-divide/src/test/java/org/apache/shenyu/plugin/divide/DividePluginTest.java b/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-divide/src/test/java/org/apache/shenyu/plugin/divide/DividePluginTest.java index 70109d02dcf3..1033b59f0a2a 100644 --- a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-divide/src/test/java/org/apache/shenyu/plugin/divide/DividePluginTest.java +++ b/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 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. */