若依 2.0

ruoyi-gateway/src/main/java/com/ruoyi/gateway/handler/SentinelFallbackHandler.java

@@ -0,0 +1,47 @@
+package com.ruoyi.gateway.handler;
+
+import java.nio.charset.StandardCharsets;
+import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.GatewayCallbackManager;
+import com.alibaba.csp.sentinel.slots.block.BlockException;
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.web.reactive.function.server.ServerResponse;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.WebExceptionHandler;
+import reactor.core.publisher.Mono;
+
+/**
+ * 自定义限流异常处理
+ * 
+ * @author ruoyi
+ */
+public class SentinelFallbackHandler implements WebExceptionHandler
+{
+    private Mono<Void> writeResponse(ServerResponse response, ServerWebExchange exchange)
+    {
+        ServerHttpResponse serverHttpResponse = exchange.getResponse();
+        serverHttpResponse.getHeaders().add("Content-Type", "application/json;charset=UTF-8");
+        byte[] datas = "{\"status\":429,\"message\":\"请求超过最大数，请稍后再试\"}".getBytes(StandardCharsets.UTF_8);
+        DataBuffer buffer = serverHttpResponse.bufferFactory().wrap(datas);
+        return serverHttpResponse.writeWith(Mono.just(buffer));
+    }
+
+    @Override
+    public Mono<Void> handle(ServerWebExchange exchange, Throwable ex)
+    {
+        if (exchange.getResponse().isCommitted())
+        {
+            return Mono.error(ex);
+        }
+        if (!BlockException.isBlockException(ex))
+        {
+            return Mono.error(ex);
+        }
+        return handleBlockedRequest(exchange, ex).flatMap(response -> writeResponse(response, exchange));
+    }
+
+    private Mono<ServerResponse> handleBlockedRequest(ServerWebExchange exchange, Throwable throwable)
+    {
+        return GatewayCallbackManager.getBlockHandler().handleRequest(exchange, throwable);
+    }
+}