Optimize usage of Sentinel RestTemplate

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

@@ -19,10 +19,12 @@ package org.springframework.cloud.alibaba.sentinel.custom;
 import java.io.IOException;
 import java.lang.reflect.Method;
 import java.net.URI;
+import java.util.Arrays;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.cloud.alibaba.sentinel.annotation.SentinelRestTemplate;
+import org.springframework.cloud.alibaba.sentinel.rest.SentinelClientHttpResponse;
 import org.springframework.http.HttpRequest;
 import org.springframework.http.client.ClientHttpRequestExecution;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
@@ -37,7 +39,7 @@ import com.alibaba.csp.sentinel.slots.block.degrade.DegradeException;
 import com.alibaba.csp.sentinel.util.StringUtil;
 
 /**
- * Interceptor using by SentinelRestTemplate and SentinelProtectInterceptor
+ * Interceptor using by SentinelRestTemplate
  *
  * @author fangjian
  */
@@ -60,7 +62,7 @@ public class SentinelProtectInterceptor implements ClientHttpRequestInterceptor
 				+ (uri.getPort() == -1 ? 80 : uri.getPort());
 		String hostWithPathResource = hostResource + uri.getPath();
 		Entry hostEntry = null, hostWithPathEntry = null;
-		ClientHttpResponse response = null;
+		ClientHttpResponse response;
 		try {
 			ContextUtil.enter(hostWithPathResource);
 			hostWithPathEntry = SphU.entry(hostWithPathResource);
@@ -68,12 +70,15 @@ public class SentinelProtectInterceptor implements ClientHttpRequestInterceptor
 			response = execution.execute(request, body);
 		}
 		catch (BlockException e) {
-			logger.error("RestTemplate block", e);
 			try {
-				handleBlockException(e);
+				return handleBlockException(request, body, execution, e);
 			}
 			catch (Exception ex) {
-				logger.error("sentinel handle BlockException error.", e);
+				if (ex instanceof IllegalStateException) {
+					throw (IllegalStateException) ex;
+				}
+				throw new IllegalStateException(
+						"sentinel handle BlockException error: " + ex.getMessage(), ex);
 			}
 		}
 		finally {
@@ -88,14 +93,18 @@ public class SentinelProtectInterceptor implements ClientHttpRequestInterceptor
 		return response;
 	}
 
-	private void handleBlockException(BlockException ex) throws Exception {
+	private ClientHttpResponse handleBlockException(HttpRequest request, byte[] body,
-		Object[] args = new Object[] { ex };
+			ClientHttpRequestExecution execution, BlockException ex) throws Exception {
+		Object[] args = new Object[] { request, body, execution, ex };
 		// handle degrade
 		if (isDegradeFailure(ex)) {
 			Method method = extractFallbackMethod(sentinelRestTemplate.fallback(),
 					sentinelRestTemplate.fallbackClass());
 			if (method != null) {
-				method.invoke(null, args);
+				return (ClientHttpResponse) method.invoke(null, args);
+			}
+			else {
+				return new SentinelClientHttpResponse();
 			}
 		}
 		// handle block
@@ -103,7 +112,10 @@ public class SentinelProtectInterceptor implements ClientHttpRequestInterceptor
 				sentinelRestTemplate.blockHandler(),
 				sentinelRestTemplate.blockHandlerClass());
 		if (blockHandler != null) {
-			blockHandler.invoke(null, args);
+			return (ClientHttpResponse) blockHandler.invoke(null, args);
+		}
+		else {
+			return new SentinelClientHttpResponse();
 		}
 	}
 
@@ -112,10 +124,25 @@ public class SentinelProtectInterceptor implements ClientHttpRequestInterceptor
 			return null;
 		}
 		Method cachedMethod = BlockClassRegistry.lookupFallback(fallbackClass, fallback);
+		Class[] args = new Class[] { HttpRequest.class, byte[].class,
+				ClientHttpRequestExecution.class, BlockException.class };
 		if (cachedMethod == null) {
-			cachedMethod = ClassUtils.getStaticMethod(fallbackClass, fallback,
+			cachedMethod = ClassUtils.getStaticMethod(fallbackClass, fallback, args);
-					BlockException.class);
+			if (cachedMethod != null) {
-			BlockClassRegistry.updateFallbackFor(fallbackClass, fallback, cachedMethod);
+				if (!ClientHttpResponse.class
+						.isAssignableFrom(cachedMethod.getReturnType())) {
+					throw new IllegalStateException(String.format(
+							"the return type of method [%s] in class [%s] is not ClientHttpResponse in degrade",
+							cachedMethod.getName(), fallbackClass.getCanonicalName()));
+				}
+				BlockClassRegistry.updateFallbackFor(fallbackClass, fallback,
+						cachedMethod);
+			}
+			else {
+				throw new IllegalStateException(String.format(
+						"Cannot find method [%s] in class [%s] with parameters %s in degrade",
+						fallback, fallbackClass.getCanonicalName(), Arrays.asList(args)));
+			}
 		}
 		return cachedMethod;
 	}
@@ -125,11 +152,25 @@ public class SentinelProtectInterceptor implements ClientHttpRequestInterceptor
 			return null;
 		}
 		Method cachedMethod = BlockClassRegistry.lookupBlockHandler(blockClass, block);
+		Class[] args = new Class[] { HttpRequest.class, byte[].class,
+				ClientHttpRequestExecution.class, BlockException.class };
 		if (cachedMethod == null) {
-			cachedMethod = ClassUtils.getStaticMethod(blockClass, block,
+			cachedMethod = ClassUtils.getStaticMethod(blockClass, block, args);
-					BlockException.class);
+			if (cachedMethod != null) {
+				if (!ClientHttpResponse.class
+						.isAssignableFrom(cachedMethod.getReturnType())) {
+					throw new IllegalStateException(String.format(
+							"the return type of method [%s] in class [%s] is not ClientHttpResponse in flow control",
+							cachedMethod.getName(), blockClass.getCanonicalName()));
+				}
 				BlockClassRegistry.updateBlockHandlerFor(blockClass, block, cachedMethod);
 			}
+			else {
+				throw new IllegalStateException(String.format(
+						"Cannot find method [%s] in class [%s] with parameters %s in flow control",
+						block, blockClass.getCanonicalName(), Arrays.asList(args)));
+			}
+		}
 		return cachedMethod;
 	}
 

spring-cloud-alibaba-sentinel/src/main/java/org/springframework/cloud/alibaba/sentinel/rest/SentinelClientHttpResponse.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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 org.springframework.cloud.alibaba.sentinel.rest;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.springframework.cloud.alibaba.sentinel.annotation.SentinelRestTemplate;
+import org.springframework.cloud.alibaba.sentinel.custom.SentinelProtectInterceptor;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
+import org.springframework.http.client.AbstractClientHttpResponse;
+
+/**
+ * Using by {@link SentinelRestTemplate} and {@link SentinelProtectInterceptor}
+ * @author <a href="mailto:fangjian0423@gmail.com">Jim</a>
+ */
+public class SentinelClientHttpResponse extends AbstractClientHttpResponse {
+
+	private final String BLOCK_STR = "RestTemplate request block by sentinel";
+
+	@Override
+	public int getRawStatusCode() throws IOException {
+		return HttpStatus.OK.value();
+	}
+
+	@Override
+	public String getStatusText() throws IOException {
+		return BLOCK_STR;
+	}
+
+	@Override
+	public void close() {
+		// nothing do
+	}
+
+	@Override
+	public InputStream getBody() throws IOException {
+		return new ByteArrayInputStream(BLOCK_STR.getBytes());
+	}
+
+	@Override
+	public HttpHeaders getHeaders() {
+		Map<String, List<String>> headers = new HashMap<>();
+		headers.put(HttpHeaders.CONTENT_TYPE,
+				Arrays.asList(MediaType.APPLICATION_JSON_UTF8_VALUE));
+		HttpHeaders httpHeaders = new HttpHeaders();
+		httpHeaders.putAll(headers);
+		return httpHeaders;
+	}
+}