添加日志

src/main/java/com/nanyan/securitylink/filter/FailoverGatewayFilterFactory.java

@@ -19,6 +19,7 @@ import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.client.reactive.ReactorClientHttpConnector;
 import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.stereotype.Component;
 import org.springframework.web.reactive.function.BodyInserters;
 import org.springframework.web.reactive.function.client.WebClient;
@@ -154,11 +155,29 @@ public class FailoverGatewayFilterFactory extends AbstractGatewayFilterFactory<F
                         ? authHeader.substring(7)
                         : null;
                 String sourceUri = exchange.getRequest().getPath().value();
+                String requestId;
+                if (exchange.getAttribute(REQUEST_ID) == null) {
+                    requestId = UUID.randomUUID().toString();
+                    exchange.getAttributes().put(REQUEST_ID, requestId);
+                }else {
+                    requestId = exchange.getAttributes().get(REQUEST_ID).toString();
+                }
+                ServerWebExchange modifiedExchange = exchange.mutate()
+                        .request(builder -> builder
+                                .headers(headers -> {
+                                    // 添加新header
+                                    headers.add(REQUEST_ID, requestId);
+                                    // 修改已有header
+                                    headers.putAll(exchange.getRequest().getHeaders());
+                                })
+                        )
+                        .build();
+                modifiedExchange.getAttributes().put(REQUEST_ID, requestId);
 
                 if (token != null) {
                     return tokenRouteMappingService.findMatchingUriConfig(token, sourceUri)
-                            .map(uriConfig -> cacheRequestBody(exchange).flatMap(cachedBody -> {
-                                log.info("Found matching URI config for token: {}, sourceUri: {}, request id:{}", token, sourceUri, exchange.getAttributes().get(REQUEST_ID));
+                            .map(uriConfig -> cacheRequestBody(modifiedExchange).flatMap(cachedBody -> {
+                                log.info("Found matching URI config for token: {}, sourceUri: {}, request id:{}", token, sourceUri, modifiedExchange.getAttributes().get(REQUEST_ID));
                                 String primaryUrl = uriConfig.getPrimaryUrl();
                                 CircularQueue errorStats = getOrCreateErrorStats(primaryUrl);
                                 long startTime = System.currentTimeMillis();
@@ -173,7 +192,7 @@ public class FailoverGatewayFilterFactory extends AbstractGatewayFilterFactory<F
                                             uriConfig.getFallbackHost(),
                                             uriConfig.getModel(),
                                             uriConfig.getModelKey(),
-                                            exchange,
+                                            modifiedExchange,
                                             false,
                                             null
                                     );
@@ -185,7 +204,7 @@ public class FailoverGatewayFilterFactory extends AbstractGatewayFilterFactory<F
                                         uriConfig.getPrimaryHost(),
                                         uriConfig.getModel(),
                                         uriConfig.getModelKey(),
-                                        exchange,
+                                        modifiedExchange,
                                         true,
                                         errorStats
                                 ).onErrorResume(throwable -> {
@@ -194,7 +213,7 @@ public class FailoverGatewayFilterFactory extends AbstractGatewayFilterFactory<F
                                                     throwable.getMessage());
                                             long duration = System.currentTimeMillis() - startTime;
                                             log.info("Source request error:{} took:{}ms, switching to target: {} request id:{}",
-                                                    throwable.getMessage(), duration, uriConfig.getTargetUri(), exchange.getAttributes().get(REQUEST_ID));
+                                                    throwable.getMessage(), duration, uriConfig.getTargetUri(), modifiedExchange.getAttributes().get(REQUEST_ID));
                                             // 记录错误
                                             getOrCreateErrorStats(uriConfig.getPrimaryUrl()).add(false);
 
@@ -205,20 +224,20 @@ public class FailoverGatewayFilterFactory extends AbstractGatewayFilterFactory<F
                                                     uriConfig.getFallbackHost(),
                                                     uriConfig.getModel(),
                                                     uriConfig.getModelKey(),
-                                                    exchange,
+                                                    modifiedExchange,
                                                     false,
                                                     null
                                             );
                                         });
                             }))
-                            .orElse(chain.filter(exchange));
+                            .orElse(chain.filter(modifiedExchange));
                 }
 
-                return chain.filter(exchange);
+                return chain.filter(modifiedExchange);
             } finally {
                 // 确保在请求处理完成后清理 ThreadLocal
                 // 清理缓存的请求体
-                exchange.getAttributes().remove(CACHED_REQUEST_BODY_KEY);
+//                exchange.getAttributes().remove(CACHED_REQUEST_BODY_KEY);
             }
         };
     }
@@ -231,12 +250,6 @@ public class FailoverGatewayFilterFactory extends AbstractGatewayFilterFactory<F
             return Mono.just("");
         }
 
-        if (exchange.getAttribute(REQUEST_ID) == null) {
-            String requestId = UUID.randomUUID().toString();
-            exchange.getAttributes().put(REQUEST_ID, requestId);
-            exchange.getRequest().getHeaders().add(REQUEST_ID, requestId);
-        }
-
         // 对于不需要请求体的方法，直接返回
         if (!requiresRequestBody(request.getMethod())) {
             return Mono.just("");
@@ -374,6 +387,7 @@ public class FailoverGatewayFilterFactory extends AbstractGatewayFilterFactory<F
 
                             exchange.getResponse().setStatusCode(responseStatus);
                             exchange.getResponse().getHeaders().putAll(clientResponse.headers().asHttpHeaders());
+                            exchange.getResponse().getHeaders().add(REQUEST_ID, exchange.getAttributes().get(REQUEST_ID).toString());
                             DataBuffer buffer = exchange.getResponse().bufferFactory()
                                     .wrap(body.getBytes(StandardCharsets.UTF_8));
                             return exchange.getResponse().writeWith(Mono.just(buffer));