Support ribbon service-level rule customization (#478)

Co-authored-by: lepdou <ledouzhang@tencent.com>

CHANGELOG.md

@@ -1,4 +1,5 @@
 # Change Log
 ---
 
+- [Feature: support ribbon service-level rule customization](https://github.com/Tencent/spring-cloud-tencent/pull/478)
 - [Feature: delete implement ServiceInstance](https://github.com/Tencent/spring-cloud-tencent/pull/481)

spring-cloud-starter-tencent-polaris-router/src/main/java/com/tencent/cloud/polaris/router/PolarisLoadBalancerCompositeRule.java

@@ -85,14 +85,21 @@ public class PolarisLoadBalancerCompositeRule extends AbstractLoadBalancerRule {
 			PolarisLoadBalancerProperties polarisLoadBalancerProperties,
 			IClientConfig iClientConfig,
 			List<RouterRequestInterceptor> requestInterceptors,
-			List<RouterResponseInterceptor> responseInterceptors) {
+			List<RouterResponseInterceptor> responseInterceptors,
+			AbstractLoadBalancerRule delegate) {
 		this.routerAPI = routerAPI;
 		this.loadBalancerProperties = polarisLoadBalancerProperties;
 		this.requestInterceptors = requestInterceptors;
 		this.responseInterceptors = responseInterceptors;
 
-		delegateRule = getRule();
+		AbstractLoadBalancerRule loadBalancerRule = getRule();
-		delegateRule.initWithNiwsConfig(iClientConfig);
+		if (loadBalancerRule != null) {
+			delegateRule = loadBalancerRule;
+			delegateRule.initWithNiwsConfig(iClientConfig);
+		}
+		else {
+			delegateRule = delegate;
+		}
 	}
 
 	@Override
@@ -176,7 +183,7 @@ public class PolarisLoadBalancerCompositeRule extends AbstractLoadBalancerRule {
 	public AbstractLoadBalancerRule getRule() {
 		String loadBalanceStrategy = loadBalancerProperties.getStrategy();
 		if (StringUtils.isEmpty(loadBalanceStrategy)) {
-			return new ZoneAvoidanceRule();
+			return null;
 		}
 		switch (loadBalanceStrategy) {
 		case STRATEGY_RANDOM:
@@ -198,4 +205,8 @@ public class PolarisLoadBalancerCompositeRule extends AbstractLoadBalancerRule {
 			return new ZoneAvoidanceRule();
 		}
 	}
+
+	public AbstractLoadBalancerRule getDelegateRule() {
+		return delegateRule;
+	}
 }

spring-cloud-starter-tencent-polaris-router/src/main/java/com/tencent/cloud/polaris/router/beanprocessor/PolarisLoadBalancerCompositeRuleBeanPostProcessor.java

@@ -0,0 +1,65 @@
+/*
+ * Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
+ *
+ * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
+ *
+ * Licensed under the BSD 3-Clause License (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://opensource.org/licenses/BSD-3-Clause
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed
+ * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+ * CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ *
+ */
+
+package com.tencent.cloud.polaris.router.beanprocessor;
+
+import java.util.List;
+
+import com.netflix.client.config.IClientConfig;
+import com.netflix.loadbalancer.AbstractLoadBalancerRule;
+import com.tencent.cloud.common.util.BeanFactoryUtils;
+import com.tencent.cloud.polaris.loadbalancer.config.PolarisLoadBalancerProperties;
+import com.tencent.cloud.polaris.router.PolarisLoadBalancerCompositeRule;
+import com.tencent.cloud.polaris.router.spi.RouterRequestInterceptor;
+import com.tencent.cloud.polaris.router.spi.RouterResponseInterceptor;
+import com.tencent.polaris.router.api.core.RouterAPI;
+
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.BeanFactoryAware;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+
+/**
+ * Decorate IRule with PolarisLoadBalancerCompositeRule.
+ *
+ * @author derekyi 2022-08-01
+ */
+public class PolarisLoadBalancerCompositeRuleBeanPostProcessor implements BeanPostProcessor, BeanFactoryAware {
+
+	private BeanFactory beanFactory;
+
+	@Override
+	public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
+		if (bean instanceof AbstractLoadBalancerRule) {
+			RouterAPI routerAPI = beanFactory.getBean(RouterAPI.class);
+			PolarisLoadBalancerProperties polarisLoadBalancerProperties = beanFactory.getBean(PolarisLoadBalancerProperties.class);
+			IClientConfig iClientConfig = beanFactory.getBean(IClientConfig.class);
+			List<RouterRequestInterceptor> requestInterceptors = BeanFactoryUtils.getBeans(beanFactory, RouterRequestInterceptor.class);
+			List<RouterResponseInterceptor> responseInterceptors = BeanFactoryUtils.getBeans(beanFactory, RouterResponseInterceptor.class);
+			return new PolarisLoadBalancerCompositeRule(routerAPI, polarisLoadBalancerProperties, iClientConfig,
+					requestInterceptors, responseInterceptors, ((AbstractLoadBalancerRule) bean));
+		}
+
+		return bean;
+	}
+
+	@Override
+	public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
+		this.beanFactory = beanFactory;
+	}
+}

spring-cloud-starter-tencent-polaris-router/src/main/java/com/tencent/cloud/polaris/router/config/RibbonConfiguration.java

@@ -18,15 +18,7 @@
 
 package com.tencent.cloud.polaris.router.config;
 
-import java.util.List;
+import com.tencent.cloud.polaris.router.beanprocessor.PolarisLoadBalancerCompositeRuleBeanPostProcessor;
-
-import com.netflix.client.config.IClientConfig;
-import com.netflix.loadbalancer.IRule;
-import com.tencent.cloud.polaris.loadbalancer.config.PolarisLoadBalancerProperties;
-import com.tencent.cloud.polaris.router.PolarisLoadBalancerCompositeRule;
-import com.tencent.cloud.polaris.router.spi.RouterRequestInterceptor;
-import com.tencent.cloud.polaris.router.spi.RouterResponseInterceptor;
-import com.tencent.polaris.router.api.core.RouterAPI;
 
 import org.springframework.context.annotation.Bean;
 
@@ -38,11 +30,7 @@ import org.springframework.context.annotation.Bean;
 public class RibbonConfiguration {
 
 	@Bean
-	public IRule polarisLoadBalancerCompositeRule(RouterAPI routerAPI,
+	public PolarisLoadBalancerCompositeRuleBeanPostProcessor polarisLoadBalancerCompositeRuleBeanPostProcessor() {
-			PolarisLoadBalancerProperties polarisLoadBalancerProperties,
+		return new PolarisLoadBalancerCompositeRuleBeanPostProcessor();
-			IClientConfig iClientConfig, List<RouterRequestInterceptor> requestInterceptors,
-			List<RouterResponseInterceptor> responseInterceptors) {
-		return new PolarisLoadBalancerCompositeRule(routerAPI, polarisLoadBalancerProperties, iClientConfig,
-				requestInterceptors, responseInterceptors);
 	}
 }

spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/PolarisLoadBalancerCompositeRuleTest.java

@@ -35,7 +35,6 @@ import com.netflix.loadbalancer.RetryRule;
 import com.netflix.loadbalancer.RoundRobinRule;
 import com.netflix.loadbalancer.Server;
 import com.netflix.loadbalancer.WeightedResponseTimeRule;
-import com.netflix.loadbalancer.ZoneAvoidanceRule;
 import com.tencent.cloud.common.metadata.MetadataContext;
 import com.tencent.cloud.common.metadata.MetadataContextHolder;
 import com.tencent.cloud.common.pojo.PolarisServer;
@@ -113,18 +112,18 @@ public class PolarisLoadBalancerCompositeRuleTest {
 	public void testGetDefaultLB() {
 		when(polarisLoadBalancerProperties.getStrategy()).thenReturn("");
 		PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-				polarisLoadBalancerProperties, config, requestInterceptors, null);
+				polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 		AbstractLoadBalancerRule defaultRule = compositeRule.getRule();
 
-		Assert.assertTrue(defaultRule instanceof ZoneAvoidanceRule);
+		Assert.assertNull(defaultRule);
 	}
 
 	@Test
 	public void testRandomLB() {
 		when(polarisLoadBalancerProperties.getStrategy()).thenReturn(PolarisLoadBalancerCompositeRule.STRATEGY_RANDOM);
 		PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-				polarisLoadBalancerProperties, config, requestInterceptors, null);
+				polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 		AbstractLoadBalancerRule lbRule = compositeRule.getRule();
 
@@ -135,7 +134,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 	public void testWeightLB() {
 		when(polarisLoadBalancerProperties.getStrategy()).thenReturn(PolarisLoadBalancerCompositeRule.STRATEGY_WEIGHT);
 		PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-				polarisLoadBalancerProperties, config, requestInterceptors, null);
+				polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 		AbstractLoadBalancerRule lbRule = compositeRule.getRule();
 
@@ -146,7 +145,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 	public void testRetryLB() {
 		when(polarisLoadBalancerProperties.getStrategy()).thenReturn(PolarisLoadBalancerCompositeRule.STRATEGY_RETRY);
 		PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-				polarisLoadBalancerProperties, config, requestInterceptors, null);
+				polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 		AbstractLoadBalancerRule lbRule = compositeRule.getRule();
 
@@ -157,7 +156,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 	public void testWeightedResponseTimeLB() {
 		when(polarisLoadBalancerProperties.getStrategy()).thenReturn(PolarisLoadBalancerCompositeRule.STRATEGY_RESPONSE_TIME_WEIGHTED);
 		PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-				polarisLoadBalancerProperties, config, requestInterceptors, null);
+				polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 		AbstractLoadBalancerRule lbRule = compositeRule.getRule();
 
@@ -168,7 +167,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 	public void tesBestAvailableLB() {
 		when(polarisLoadBalancerProperties.getStrategy()).thenReturn(PolarisLoadBalancerCompositeRule.STRATEGY_BEST_AVAILABLE);
 		PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-				polarisLoadBalancerProperties, config, requestInterceptors, null);
+				polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 		AbstractLoadBalancerRule lbRule = compositeRule.getRule();
 
@@ -179,7 +178,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 	public void tesRoundRobinLB() {
 		when(polarisLoadBalancerProperties.getStrategy()).thenReturn(PolarisLoadBalancerCompositeRule.STRATEGY_ROUND_ROBIN);
 		PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-				polarisLoadBalancerProperties, config, requestInterceptors, null);
+				polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 		AbstractLoadBalancerRule lbRule = compositeRule.getRule();
 
@@ -190,7 +189,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 	public void testAvailabilityFilteringLB() {
 		when(polarisLoadBalancerProperties.getStrategy()).thenReturn(PolarisLoadBalancerCompositeRule.STRATEGY_AVAILABILITY_FILTERING);
 		PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-				polarisLoadBalancerProperties, config, requestInterceptors, null);
+				polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 		AbstractLoadBalancerRule lbRule = compositeRule.getRule();
 
@@ -209,7 +208,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 			setTransitiveMetadata();
 
 			PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-					polarisLoadBalancerProperties, config, requestInterceptors, null);
+					polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 			ServiceInstances serviceInstances = assembleServiceInstances();
 			PolarisRouterContext routerContext = assembleRouterContext();
@@ -244,7 +243,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 			setTransitiveMetadata();
 
 			PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-					polarisLoadBalancerProperties, config, requestInterceptors, null);
+					polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 			ServiceInstances serviceInstances = assembleServiceInstances();
 			PolarisRouterContext routerContext = assembleRouterContext();
@@ -275,7 +274,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 			setTransitiveMetadata();
 
 			PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-					polarisLoadBalancerProperties, config, requestInterceptors, null);
+					polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 			ServiceInstances serviceInstances = assembleServiceInstances();
 			PolarisRouterContext routerContext = assembleRouterContext();
@@ -306,7 +305,7 @@ public class PolarisLoadBalancerCompositeRuleTest {
 			setTransitiveMetadata();
 
 			PolarisLoadBalancerCompositeRule compositeRule = new PolarisLoadBalancerCompositeRule(routerAPI,
-					polarisLoadBalancerProperties, config, requestInterceptors, null);
+					polarisLoadBalancerProperties, config, requestInterceptors, null, null);
 
 			ProcessRoutersResponse assembleResponse = assembleProcessRoutersResponse();
 			when(routerAPI.processRouters(any())).thenReturn(assembleResponse);

spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/beanprocessor/PolarisLoadBalancerCompositeRuleBeanPostProcessorTest.java

@@ -0,0 +1,148 @@
+package com.tencent.cloud.polaris.router.beanprocessor;
+
+import com.netflix.loadbalancer.AbstractLoadBalancerRule;
+import com.netflix.loadbalancer.BestAvailableRule;
+import com.netflix.loadbalancer.IRule;
+import com.netflix.loadbalancer.RandomRule;
+import com.netflix.loadbalancer.RoundRobinRule;
+import com.netflix.loadbalancer.ZoneAvoidanceRule;
+import com.tencent.cloud.polaris.loadbalancer.config.PolarisLoadBalancerProperties;
+import com.tencent.cloud.polaris.router.PolarisLoadBalancerCompositeRule;
+import com.tencent.cloud.polaris.router.config.RibbonConfiguration;
+import com.tencent.polaris.client.api.SDKContext;
+import com.tencent.polaris.router.api.core.RouterAPI;
+import com.tencent.polaris.router.client.api.DefaultRouterAPI;
+import org.junit.Assert;
+import org.junit.Test;
+
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.context.properties.EnableConfigurationProperties;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.cloud.netflix.ribbon.RibbonAutoConfiguration;
+import org.springframework.cloud.netflix.ribbon.RibbonClient;
+import org.springframework.cloud.netflix.ribbon.RibbonClients;
+import org.springframework.cloud.netflix.ribbon.SpringClientFactory;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+
+/**
+ * Test for {@link PolarisLoadBalancerCompositeRuleBeanPostProcessor}.
+ *
+ * @author derekyi 2022-08-01
+ */
+public class PolarisLoadBalancerCompositeRuleBeanPostProcessorTest {
+
+	private static final String SERVICE_1 = "service1";
+
+	private static final String SERVICE_2 = "service2";
+
+	@Test
+	public void test1() {
+		ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+				.withConfiguration(AutoConfigurations.of(RibbonDefaultConfig.class, PolarisRibbonTest.class, RibbonAutoConfiguration.class));
+		contextRunner.run(context -> {
+			SpringClientFactory springClientFactory = context.getBean(SpringClientFactory.class);
+
+			IRule rule = springClientFactory.getInstance(SERVICE_1, IRule.class);
+			Assert.assertTrue(rule instanceof PolarisLoadBalancerCompositeRule);
+			AbstractLoadBalancerRule delegateRule = ((PolarisLoadBalancerCompositeRule) rule).getDelegateRule();
+			//ZoneAvoidanceRule default
+			Assert.assertTrue(delegateRule instanceof ZoneAvoidanceRule);
+		});
+	}
+
+	@Test
+	public void test2() {
+		ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+				.withConfiguration(AutoConfigurations.of(RibbonDefaultConfig.class, PolarisRibbonTest.class, RibbonAutoConfiguration.class))
+				.withPropertyValues("spring.cloud.polaris.loadbalancer.strategy = random");
+		contextRunner.run(context -> {
+			SpringClientFactory springClientFactory = context.getBean(SpringClientFactory.class);
+
+			IRule rule = springClientFactory.getInstance(SERVICE_1, IRule.class);
+			Assert.assertTrue(rule instanceof PolarisLoadBalancerCompositeRule);
+			AbstractLoadBalancerRule delegateRule = ((PolarisLoadBalancerCompositeRule) rule).getDelegateRule();
+			//spring.cloud.polaris.loadbalancer.strategy = random
+			Assert.assertTrue(delegateRule instanceof RandomRule);
+		});
+	}
+
+	@Test
+	public void test3() {
+		ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+				.withConfiguration(AutoConfigurations.of(RibbonDefaultConfig.class, PolarisRibbonTest.class, RibbonAutoConfiguration.class))
+				.withPropertyValues("service1.ribbon.NFLoadBalancerRuleClassName=com.netflix.loadbalancer.RoundRobinRule");
+		contextRunner.run(context -> {
+			SpringClientFactory springClientFactory = context.getBean(SpringClientFactory.class);
+
+			IRule rule1 = springClientFactory.getInstance(SERVICE_1, IRule.class);
+			Assert.assertTrue(rule1 instanceof PolarisLoadBalancerCompositeRule);
+			AbstractLoadBalancerRule delegateRule1 = ((PolarisLoadBalancerCompositeRule) rule1).getDelegateRule();
+			//service1.ribbon.NFLoadBalancerRuleClassName=com.netflix.loadbalancer.RoundRobinRule
+			Assert.assertTrue(delegateRule1 instanceof RoundRobinRule);
+
+			IRule rule2 = springClientFactory.getInstance(SERVICE_2, IRule.class);
+			Assert.assertTrue(rule2 instanceof PolarisLoadBalancerCompositeRule);
+
+			AbstractLoadBalancerRule delegateRule2 = ((PolarisLoadBalancerCompositeRule) rule2).getDelegateRule();
+			//ZoneAvoidanceRule default
+			Assert.assertTrue(delegateRule2 instanceof ZoneAvoidanceRule);
+		});
+	}
+
+	@Test
+	public void test4() {
+		ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+				.withConfiguration(AutoConfigurations.of(CustomRibbonConfig.class, PolarisRibbonTest.class, RibbonAutoConfiguration.class));
+		contextRunner.run(context -> {
+			SpringClientFactory springClientFactory = context.getBean(SpringClientFactory.class);
+
+			IRule rule1 = springClientFactory.getInstance(SERVICE_1, IRule.class);
+			Assert.assertTrue(rule1 instanceof PolarisLoadBalancerCompositeRule);
+			AbstractLoadBalancerRule delegateRule1 = ((PolarisLoadBalancerCompositeRule) rule1).getDelegateRule();
+			//RibbonConfigForService1#loadBalancerRule returns BestAvailableRule
+			Assert.assertTrue(delegateRule1 instanceof BestAvailableRule);
+
+			IRule rule2 = springClientFactory.getInstance(SERVICE_2, IRule.class);
+			Assert.assertTrue(rule2 instanceof PolarisLoadBalancerCompositeRule);
+			AbstractLoadBalancerRule delegateRule2 = ((PolarisLoadBalancerCompositeRule) rule2).getDelegateRule();
+			//ZoneAvoidanceRule default
+			Assert.assertTrue(delegateRule2 instanceof ZoneAvoidanceRule);
+		});
+	}
+
+	@Configuration
+	@RibbonClients(defaultConfiguration = {RibbonConfiguration.class})
+	static class RibbonDefaultConfig {
+	}
+
+	@Configuration
+	@RibbonClients(value = {@RibbonClient(name = SERVICE_1, configuration = RibbonConfigForService1.class)}, defaultConfiguration = RibbonConfiguration.class)
+	static class CustomRibbonConfig {
+	}
+
+	static class RibbonConfigForService1 {
+
+		@Bean
+		public IRule loadBalancerRule() {
+			return new BestAvailableRule();
+		}
+	}
+
+
+	@Configuration
+	@EnableConfigurationProperties(PolarisLoadBalancerProperties.class)
+	static class PolarisRibbonTest {
+
+		@Bean
+		public SDKContext sdkContext() {
+			return SDKContext.initContext();
+		}
+
+		@Bean
+		public RouterAPI routerAPI(SDKContext sdkContext) {
+			return new DefaultRouterAPI(sdkContext);
+		}
+	}
+}