Merge pull request #864 from Rivers-Shall/fix-#762

Fix #762 SentinelRestTemplate enhancement

spring-cloud-alibaba-examples/nacos-example/nacos-discovery-example/nacos-discovery-consumer-example/src/main/java/com/alibaba/cloud/examples/ConsumerApplication.java

@@ -26,6 +26,7 @@ public class ConsumerApplication {
 
 	@LoadBalanced
 	@Bean
+	@SentinelRestTemplate(urlCleanerClass = UrlCleaner.class, urlCleaner = "clean")
 	public RestTemplate restTemplate() {
 		return new RestTemplate();
 	}

spring-cloud-alibaba-examples/nacos-example/nacos-discovery-example/nacos-discovery-consumer-example/src/main/java/com/alibaba/cloud/examples/UrlCleaner.java

@@ -0,0 +1,12 @@
+package com.alibaba.cloud.examples;
+
+public class UrlCleaner {
+    public static String clean(String url) {
+        System.out.println("enter urlCleaner");
+        if (url.matches(".*/echo/.*")) {
+            System.out.println("change url");
+            url = url.replaceAll("/echo/.*", "/echo/{str}");
+        }
+        return url;
+    }
+}

spring-cloud-alibaba-sentinel/src/main/java/com/alibaba/cloud/sentinel/SentinelConstants.java

@@ -25,6 +25,7 @@ public interface SentinelConstants {
 
 	String BLOCK_TYPE = "block";
 	String FALLBACK_TYPE = "fallback";
+	String URLCLEANER_TYPE = "urlCleaner";
 
 	// commercialization
 

spring-cloud-alibaba-sentinel/src/main/java/com/alibaba/cloud/sentinel/annotation/SentinelRestTemplate.java

@@ -38,4 +38,7 @@ public @interface SentinelRestTemplate {
 
 	Class<?> fallbackClass() default void.class;
 
+	String urlCleaner() default "";
+
+	Class<?> urlCleanerClass() default void.class;
 }

spring-cloud-alibaba-sentinel/src/main/java/com/alibaba/cloud/sentinel/custom/BlockClassRegistry.java

@@ -21,6 +21,7 @@ import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
 import com.alibaba.csp.sentinel.util.StringUtil;
+import org.springframework.util.StringUtils;
 
 /**
  * @author fangjian
@@ -29,6 +30,7 @@ final class BlockClassRegistry {
 
 	private static final Map<String, Method> FALLBACK_MAP = new ConcurrentHashMap<>();
 	private static final Map<String, Method> BLOCK_HANDLER_MAP = new ConcurrentHashMap<>();
+	private static final Map<String, Method> URL_CLEANER_MAP = new ConcurrentHashMap<>();
 
 	static Method lookupFallback(Class<?> clazz, String name) {
 		return FALLBACK_MAP.get(getKey(clazz, name));
@@ -38,6 +40,10 @@ final class BlockClassRegistry {
 		return BLOCK_HANDLER_MAP.get(getKey(clazz, name));
 	}
 
+	static Method lookupUrlCleaner(Class<?> clazz, String name) {
+		return URL_CLEANER_MAP.get(getKey(clazz, name));
+	}
+
 	static void updateFallbackFor(Class<?> clazz, String name, Method method) {
 		if (clazz == null || StringUtil.isBlank(name)) {
 			throw new IllegalArgumentException("Bad argument");
@@ -52,6 +58,13 @@ final class BlockClassRegistry {
 		BLOCK_HANDLER_MAP.put(getKey(clazz, name), method);
 	}
 
+	static void updateUrlCleanerFor(Class<?> clazz, String name, Method method) {
+		if (clazz == null || StringUtil.isBlank(name)) {
+			throw new IllegalArgumentException("Bad argument");
+		}
+		URL_CLEANER_MAP.put(getKey(clazz, name), method);
+	}
+
 	private static String getKey(Class<?> clazz, String name) {
 		return String.format("%s:%s", clazz.getCanonicalName(), name);
 	}

spring-cloud-alibaba-sentinel/src/main/java/com/alibaba/cloud/sentinel/custom/SentinelBeanPostProcessor.java

@@ -90,6 +90,9 @@ public class SentinelBeanPostProcessor implements MergedBeanDefinitionPostProces
 		checkBlock4RestTemplate(sentinelRestTemplate.fallbackClass(),
 				sentinelRestTemplate.fallback(), beanName,
 				SentinelConstants.FALLBACK_TYPE);
+		checkBlock4RestTemplate(sentinelRestTemplate.urlCleanerClass(),
+				sentinelRestTemplate.urlCleaner(), beanName,
+				SentinelConstants.URLCLEANER_TYPE);
 	}
 
 	private void checkBlock4RestTemplate(Class<?> blockClass, String blockMethod,
@@ -111,8 +114,14 @@ public class SentinelBeanPostProcessor implements MergedBeanDefinitionPostProces
 			throw new IllegalArgumentException(type + " method attribute exists but "
 					+ type + " class attribute is not exists in bean[" + beanName + "]");
 		}
-		Class[] args = new Class[] { HttpRequest.class, byte[].class,
+		Class[] args;
+		if (type.equals(SentinelConstants.URLCLEANER_TYPE)) {
+			args = new Class[] { String.class };
+		}
+		else {
+			args = new Class[] { HttpRequest.class, byte[].class,
 					ClientHttpRequestExecution.class, BlockException.class };
+		}
 		String argsStr = Arrays.toString(
 				Arrays.stream(args).map(clazz -> clazz.getSimpleName()).toArray());
 		Method foundMethod = ClassUtils.getStaticMethod(blockClass, blockMethod, args);
@@ -127,21 +136,32 @@ public class SentinelBeanPostProcessor implements MergedBeanDefinitionPostProces
 					+ ", please check your class name, method name and arguments");
 		}
 
-		if (!ClientHttpResponse.class.isAssignableFrom(foundMethod.getReturnType())) {
+		Class<?> standardReturnType;
-			log.error(
+		if (type.equals(SentinelConstants.URLCLEANER_TYPE)) {
-					"{} method return value in bean[{}] is not ClientHttpResponse: {}#{}{}",
+			standardReturnType = String.class;
-					type, beanName, blockClass.getName(), blockMethod, argsStr);
+		}
+		else {
+			standardReturnType = ClientHttpResponse.class;
+		}
+
+		if (!standardReturnType.isAssignableFrom(foundMethod.getReturnType())) {
+			log.error("{} method return value in bean[{}] is not {}: {}#{}{}", type,
+					beanName, standardReturnType.getName(), blockClass.getName(),
+					blockMethod, argsStr);
 			throw new IllegalArgumentException(type + " method return value in bean["
-					+ beanName + "] is not ClientHttpResponse: " + blockClass.getName()
+					+ beanName + "] is not " + standardReturnType.getName() + ": "
-					+ "#" + blockMethod + argsStr);
+					+ blockClass.getName() + "#" + blockMethod + argsStr);
 		}
 		if (type.equals(SentinelConstants.BLOCK_TYPE)) {
 			BlockClassRegistry.updateBlockHandlerFor(blockClass, blockMethod,
 					foundMethod);
 		}
-		else {
+		else if (type.equals(SentinelConstants.FALLBACK_TYPE)) {
 			BlockClassRegistry.updateFallbackFor(blockClass, blockMethod, foundMethod);
 		}
+		else {
+			BlockClassRegistry.updateUrlCleanerFor(blockClass, blockMethod, foundMethod);
+		}
 	}
 
 	private boolean checkSentinelProtect(RootBeanDefinition beanDefinition,
@@ -177,7 +197,9 @@ public class SentinelBeanPostProcessor implements MergedBeanDefinitionPostProces
 					.append(sentinelRestTemplate.blockHandlerClass().getSimpleName())
 					.append(sentinelRestTemplate.blockHandler()).append("_")
 					.append(sentinelRestTemplate.fallbackClass().getSimpleName())
-					.append(sentinelRestTemplate.fallback());
+					.append(sentinelRestTemplate.fallback()).append("_")
+					.append(sentinelRestTemplate.urlCleanerClass().getSimpleName())
+					.append(sentinelRestTemplate.urlCleaner());
 			RestTemplate restTemplate = (RestTemplate) bean;
 			String interceptorBeanName = interceptorBeanNamePrefix + "@"
 					+ bean.toString();

spring-cloud-alibaba-sentinel/src/main/java/com/alibaba/cloud/sentinel/custom/SentinelProtectInterceptor.java

@@ -65,6 +65,14 @@ public class SentinelProtectInterceptor implements ClientHttpRequestInterceptor
 		if (hostResource.equals(hostWithPathResource)) {
 			entryWithPath = false;
 		}
+		Method urlCleanerMethod = BlockClassRegistry.lookupUrlCleaner(
+				sentinelRestTemplate.urlCleanerClass(),
+				sentinelRestTemplate.urlCleaner());
+		if (urlCleanerMethod != null) {
+			hostWithPathResource = (String) methodInvoke(urlCleanerMethod,
+					hostWithPathResource);
+		}
+
 		Entry hostEntry = null, hostWithPathEntry = null;
 		ClientHttpResponse response = null;
 		try {
@@ -105,7 +113,7 @@ public class SentinelProtectInterceptor implements ClientHttpRequestInterceptor
 			Method fallbackMethod = extractFallbackMethod(sentinelRestTemplate.fallback(),
 					sentinelRestTemplate.fallbackClass());
 			if (fallbackMethod != null) {
-				return methodInvoke(fallbackMethod, args);
+				return (ClientHttpResponse) methodInvoke(fallbackMethod, args);
 			}
 			else {
 				return new SentinelClientHttpResponse();
@@ -116,16 +124,16 @@ public class SentinelProtectInterceptor implements ClientHttpRequestInterceptor
 				sentinelRestTemplate.blockHandler(),
 				sentinelRestTemplate.blockHandlerClass());
 		if (blockHandler != null) {
-			return methodInvoke(blockHandler, args);
+			return (ClientHttpResponse) methodInvoke(blockHandler, args);
 		}
 		else {
 			return new SentinelClientHttpResponse();
 		}
 	}
 
-	private ClientHttpResponse methodInvoke(Method method, Object... args) {
+	private Object methodInvoke(Method method, Object... args) {
 		try {
-			return (ClientHttpResponse) method.invoke(null, args);
+			return method.invoke(null, args);
 		}
 		catch (IllegalAccessException e) {
 			throw new RuntimeException(e);

spring-cloud-alibaba-sentinel/src/test/java/com/alibaba/cloud/sentinel/SentinelRestTemplateTests.java

@@ -91,6 +91,26 @@ public class SentinelRestTemplateTests {
 		new AnnotationConfigApplicationContext(TestConfig10.class);
 	}
 
+	@Test(expected = BeanCreationException.class)
+	public void testUrlClnMethod() {
+		new AnnotationConfigApplicationContext(TestConfig11.class);
+	}
+
+	@Test(expected = BeanCreationException.class)
+	public void testUrlClnClass() {
+		new AnnotationConfigApplicationContext(TestConfig12.class);
+	}
+
+	@Test(expected = BeanCreationException.class)
+	public void testUrlClnMethodExists() {
+		new AnnotationConfigApplicationContext(TestConfig13.class);
+	}
+
+	@Test(expected = BeanCreationException.class)
+	public void testUrlClnReturnValue() {
+		new AnnotationConfigApplicationContext(TestConfig14.class);
+	}
+
 	@Configuration
 	public static class TestConfig1 {
 		@Bean
@@ -160,7 +180,7 @@ public class SentinelRestTemplateTests {
 		}
 
 		@Bean
-		@SentinelRestTemplate(blockHandlerClass = SentinelRestTemplateTests.ExceptionUtil.class, blockHandler = "handleException", fallbackClass = SentinelRestTemplateTests.ExceptionUtil.class, fallback = "fallbackException")
+		@SentinelRestTemplate(blockHandlerClass = SentinelRestTemplateTests.ExceptionUtil.class, blockHandler = "handleException", fallbackClass = SentinelRestTemplateTests.ExceptionUtil.class, fallback = "fallbackException", urlCleanerClass = SentinelRestTemplateTests.UrlCleanUtil.class, urlCleaner = "clean")
 		RestTemplate restTemplate() {
 			return new RestTemplate();
 		}
@@ -247,6 +267,66 @@ public class SentinelRestTemplateTests {
 		}
 	}
 
+	@Configuration
+	public static class TestConfig11 {
+		@Bean
+		SentinelBeanPostProcessor sentinelBeanPostProcessor(
+				ApplicationContext applicationContext) {
+			return new SentinelBeanPostProcessor(applicationContext);
+		}
+
+		@Bean
+		@SentinelRestTemplate(urlCleaner = "cln")
+		RestTemplate restTemplate() {
+			return new RestTemplate();
+		}
+	}
+
+	@Configuration
+	public static class TestConfig12 {
+		@Bean
+		SentinelBeanPostProcessor sentinelBeanPostProcessor(
+				ApplicationContext applicationContext) {
+			return new SentinelBeanPostProcessor(applicationContext);
+		}
+
+		@Bean
+		@SentinelRestTemplate(urlCleanerClass = UrlCleanUtil.class)
+		RestTemplate restTemplate() {
+			return new RestTemplate();
+		}
+	}
+
+	@Configuration
+	public static class TestConfig13 {
+		@Bean
+		SentinelBeanPostProcessor sentinelBeanPostProcessor(
+				ApplicationContext applicationContext) {
+			return new SentinelBeanPostProcessor(applicationContext);
+		}
+
+		@Bean
+		@SentinelRestTemplate(urlCleanerClass = SentinelRestTemplateTests.UrlCleanUtil.class, urlCleaner = "clean1")
+		RestTemplate restTemplate() {
+			return new RestTemplate();
+		}
+	}
+
+	@Configuration
+	public static class TestConfig14 {
+		@Bean
+		SentinelBeanPostProcessor sentinelBeanPostProcessor(
+				ApplicationContext applicationContext) {
+			return new SentinelBeanPostProcessor(applicationContext);
+		}
+
+		@Bean
+		@SentinelRestTemplate(urlCleanerClass = SentinelRestTemplateTests.UrlCleanUtil.class, urlCleaner = "clean2")
+		RestTemplate restTemplate() {
+			return new RestTemplate();
+		}
+	}
+
 	public static class ExceptionUtil {
 		public static SentinelClientHttpResponse handleException(HttpRequest request,
 				byte[] body, ClientHttpRequestExecution execution, BlockException ex) {
@@ -271,4 +351,13 @@ public class SentinelRestTemplateTests {
 		}
 	}
 
+	public static class UrlCleanUtil {
+		public static String clean(String url) {
+			return url;
+		}
+
+		public static void clean2(String url) {
+		}
+	}
+
 }