Spring Cloud Gateway 流量复制

1
import java.net.URI;
2
import java.net.URLDecoder;
3
import java.net.URLEncoder;
4
import java.nio.charset.StandardCharsets;
5
import java.time.Duration;
6
import java.util.List;
7
import java.util.concurrent.ExecutorService;
8
import java.util.concurrent.Executors;
9

10
import org.springframework.beans.factory.DisposableBean;
11
import org.springframework.cloud.gateway.filter.GatewayFilter;
12
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
13
import org.springframework.core.io.buffer.DataBuffer;
14
import org.springframework.core.io.buffer.DataBufferUtils;
15
import org.springframework.http.HttpMethod;
16
import org.springframework.http.HttpStatus;
17
import org.springframework.http.server.reactive.ServerHttpRequest;
18
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
19
import org.springframework.stereotype.Component;
20
import org.springframework.util.CollectionUtils;
21
import org.springframework.web.reactive.function.BodyInserters;
22
import org.springframework.web.reactive.function.client.WebClient;
23
import org.springframework.web.reactive.function.client.WebClientResponseException;
24

25
import lombok.Data;
26
import lombok.extern.slf4j.Slf4j;
27
import reactor.core.publisher.Flux;
28
import reactor.core.publisher.Mono;
29
import reactor.core.scheduler.Scheduler;
30
import reactor.core.scheduler.Schedulers;
31

32
/**
33
 * 流量镜像网关过滤器工厂
34
 *
35
 * 功能：将请求复制到指定的镜像路由，实现完全非阻塞流量镜像 特性： - 真正的异步非阻塞处理，镜像不影响主请求性能 - 支持流式请求体处理，避免大请求体内存占用 - 支持多个镜像目标 - 自动错误处理和资源清理 - 智能的请求体复用机制
36
 */
37
@Slf4j
38
@Component
39
public class TrafficMirroringGatewayFilterFactory
40
    extends AbstractGatewayFilterFactory<TrafficMirroringGatewayFilterFactory.Config> implements DisposableBean {
41

42
    private final WebClient webClient;
43
    private final ExecutorService executorService;
44
    private final Scheduler scheduler;
45

46
    public TrafficMirroringGatewayFilterFactory() {
47
        super(TrafficMirroringGatewayFilterFactory.Config.class);
48
        this.executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2, r -> {
49
            Thread t = new Thread(r, "traffic-mirror-");
50
            t.setDaemon(true);
51
            return t;
52
        });
53
        this.scheduler = Schedulers.fromExecutorService(executorService);
54
        this.webClient =
55
            WebClient.builder().codecs(configurer -> configurer.defaultCodecs().maxInMemorySize(1024 * 1024)) // 1MB
56
                .build();
57
    }
58

59
    @Override
60
    public GatewayFilter apply(Config config) {
61
        return (exchange, chain) -> {
62
            List<URI> mirroredRoutes = config.getMirroredRoutes();
63
            if (CollectionUtils.isEmpty(mirroredRoutes)) {
64
                return chain.filter(exchange);
65
            }
66

67
            ServerHttpRequest request = exchange.getRequest();
68

69
            // 对于无请求体的方法，立即镜像并继续主请求
70
            if (isBodylessMethod(request.getMethod())) {
71
                // 完全异步镜像，不阻塞主请求
72
                Mono.fromRunnable(() -> mirrorBodylessRequest(request, mirroredRoutes, config)).subscribeOn(scheduler)
73
                    .doOnError(error -> log.warn("Bodyless mirror failed", error)).subscribe();
74

75
                return chain.filter(exchange);
76
            }
77

78
            // 对于有请求体的方法，使用流式处理，不等待请求体完全读取
79
            ServerHttpRequest decoratedRequest = decorateRequestForMirroring(request, mirroredRoutes, config);
80

81
            // 立即继续主请求，镜像在后台异步进行
82
            return chain.filter(exchange.mutate().request(decoratedRequest).build());
83
        };
84
    }
85

86
    /**
87
     * 装饰请求以支持流式镜像，不阻塞主请求
88
     */
89
    private ServerHttpRequest decorateRequestForMirroring(ServerHttpRequest originalRequest, List<URI> mirroredRoutes,
90
        Config config) {
91
        return new ServerHttpRequestDecorator(originalRequest) {
92
            @Override
93
            public Flux<DataBuffer> getBody() {
94
                Flux<DataBuffer> originalBody = super.getBody();
95

96
                // 收集完整的请求体，然后进行镜像
97
                return originalBody.map(dataBuffer -> {
98
                    // 克隆DataBuffer用于镜像，保持原始DataBuffer不变
99
                    DataBuffer cloned = cloneDataBuffer(dataBuffer);
100
                    return new DataBufferWrapper(dataBuffer, cloned);
101
                }).collectList() // 收集所有DataBufferWrapper
102
                    .doOnNext(wrappers -> {
103
                        // 异步处理镜像，不阻塞主请求
104
                        Mono.fromRunnable(() -> {
105
                            try {
106
                                // 将所有克隆的DataBuffer合并为完整的请求体用于镜像
107
                                java.io.ByteArrayOutputStream outputStream = new java.io.ByteArrayOutputStream();
108
                                for (DataBufferWrapper wrapper : wrappers) {
109
                                    byte[] bytes = new byte[wrapper.cloned.readableByteCount()];
110
                                    wrapper.cloned.read(bytes);
111
                                    outputStream.write(bytes);
112
                                    // 释放克隆的DataBuffer
113
                                    DataBufferUtils.release(wrapper.cloned);
114
                                }
115
                                byte[] fullBodyBytes = outputStream.toByteArray();
116

117
                                // 发送镜像请求
118
                                mirrorRequest(originalRequest, mirroredRoutes, config, fullBodyBytes);
119
                            } catch (Exception e) {
120
                                log.warn("Failed to mirror request", e);
121
                            }
122
                        }).subscribeOn(scheduler).doOnError(error -> log.warn("Mirror request failed", error))
123
                            .subscribe();
124
                    }).doOnError(error -> log.warn("Request body collection error", error))
125
                    .flatMapMany(wrappers -> Flux.fromIterable(wrappers).map(wrapper -> wrapper.original));
126
            }
127
        };
128
    }
129

130
    /**
131
     * 克隆数据缓冲区用于镜像
132
     */
133
    private DataBuffer cloneDataBuffer(DataBuffer original) {
134
        DataBuffer cloned = original.factory().allocateBuffer(original.readableByteCount());
135
        cloned.write(original.slice(original.readPosition(), original.readableByteCount()));
136
        return cloned;
137
    }
138

139
    /**
140
     * DataBuffer包装器，用于同时持有原始和克隆的DataBuffer
141
     */
142
    private static class DataBufferWrapper {
143
        final DataBuffer original;
144
        final DataBuffer cloned;
145

146
        DataBufferWrapper(DataBuffer original, DataBuffer cloned) {
147
            this.original = original;
148
            this.cloned = cloned;
149
        }
150
    }
151

152
    /**
153
     * 判断是否为无请求体的HTTP方法
154
     */
155
    private boolean isBodylessMethod(HttpMethod method) {
156
        return method == HttpMethod.GET || method == HttpMethod.HEAD || method == HttpMethod.OPTIONS
157
            || method == HttpMethod.TRACE;
158
    }
159

160
    /**
161
     * 执行镜像请求 - 完全异步，不阻塞调用者
162
     */
163
    private void mirrorRequest(ServerHttpRequest request, List<URI> mirroredRoutes, Config config, byte[] bodyBytes) {
164
        mirroredRoutes.forEach(mirroredRoute -> {
165
            Mono.fromCallable(() -> {
166
                WebClient.RequestHeadersSpec<?> requestSpec = webClient.method(request.getMethod())
167
                    .uri(buildMirrorUri(mirroredRoute, request)).headers(httpHeaders -> {
168
                        httpHeaders.addAll(request.getHeaders());
169
                        // 移除可能导致问题的头部
170
                        httpHeaders.remove("host");
171
                        httpHeaders.remove("content-length");
172
                        if (bodyBytes != null && bodyBytes.length == 0) {
173
                            httpHeaders.remove("content-type");
174
                        }
175
                    });
176

177
                if (bodyBytes != null && bodyBytes.length > 0) {
178
                    requestSpec = ((WebClient.RequestBodySpec)requestSpec).body(BodyInserters.fromValue(bodyBytes));
179
                }
180

181
                return requestSpec.retrieve().toBodilessEntity().timeout(Duration.ofSeconds(config.getTimeoutSeconds()))
182
                    .doOnSuccess(response -> {
183
                        log.info("Mirror request succeeded to {}", mirroredRoute);
184
                        // if (log.isDebugEnabled()) {
185
                        // log.debug("Mirror request succeeded to {}", mirroredRoute);
186
                        // }
187
                    }).doOnError(error -> handleMirrorError(mirroredRoute, error, config));
188
            }).subscribeOn(scheduler).flatMap(mono -> mono)
189
                .doOnError(error -> handleMirrorError(mirroredRoute, error, config)).onErrorResume(error -> {
190
                    log.error(error.getMessage(), error);
191
                    return Mono.empty();
192
                }) // 忽略镜像错误，不影响主流程
193
                .subscribe();
194
        });
195
    }
196

197
    /**
198
     * 镜像无请求体的请求
199
     */
200
    private void mirrorBodylessRequest(ServerHttpRequest request, List<URI> mirroredRoutes, Config config) {
201
        mirrorRequest(request, mirroredRoutes, config, null);
202
    }
203

204
    /**
205
     * 构建镜像URI，保留原始请求的路径和查询参数
206
     */
207
    private URI buildMirrorUri(URI mirroredRoute, ServerHttpRequest request) {
208
        String query = request.getURI().getQuery();
209
        String uriString = mirroredRoute.toString();
210
        if (query != null && !query.isEmpty() && !isUrlEncoded(query)) {
211
            try {
212
                uriString += "?" + URLEncoder.encode(query, StandardCharsets.UTF_8.toString());
213
            } catch (Exception e) {
214
                // ignore
215
            }
216
        }
217
        return URI.create(uriString);
218
    }
219

220
    public static boolean isUrlEncoded(String value) {
221
        if (value == null)
222
            return false;
223

224
        try {
225
            // 解码后与原字符串不同，则可能已编码
226
            String decoded = URLDecoder.decode(value, StandardCharsets.UTF_8.toString());
227
            return !decoded.equals(value);
228
        } catch (Exception e) {
229
            // 解码失败说明可能未编码
230
            return false;
231
        }
232
    }
233

234
    /**
235
     * 处理镜像请求错误
236
     */
237
    private void handleMirrorError(URI mirroredRoute, Throwable error, Config config) {
238
        if (config.isEnableErrorLogging()) {
239
            if (error instanceof WebClientResponseException) {
240
                WebClientResponseException webClientError = (WebClientResponseException)error;
241
                // 404 不记录警告，可能是正常情况
242
                if (webClientError.getStatusCode() != HttpStatus.NOT_FOUND) {
243
                    log.warn("Mirror request failed to {}: {} - {}", mirroredRoute, webClientError.getStatusCode(),
244
                        webClientError.getMessage());
245
                }
246
            } else {
247
                log.warn("Mirror request failed to {}: {}", mirroredRoute, error.getMessage());
248
            }
249
        }
250
    }
251

252
    @Override
253
    public String name() {
254
        return "TrafficMirroring";
255
    }
256

257
    /**
258
     * 清理资源
259
     */
260
    @Override
261
    public void destroy() throws Exception {
262
        log.info("Shutting down traffic mirroring executor service");
263
        scheduler.dispose();
264
        executorService.shutdown();
265
        try {
266
            if (!executorService.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS)) {
267
                log.warn("Executor service did not terminate gracefully, forcing shutdown");
268
                executorService.shutdownNow();
269
            }
270
        } catch (InterruptedException e) {
271
            log.warn("Interrupted while waiting for executor service termination");
272
            executorService.shutdownNow();
273
            Thread.currentThread().interrupt();
274
        }
275
    }
276

277
    /**
278
     * 流量镜像配置类
279
     */
280
    @Data
281
    public static class Config {
282
        /**
283
         * 镜像目标路由列表
284
         */
285
        private List<URI> mirroredRoutes;
286

287
        /**
288
         * 镜像请求超时时间（秒），默认3秒
289
         */
290
        private int timeoutSeconds = 3;
291

292
        /**
293
         * 是否启用错误日志，默认启用
294
         */
295
        private boolean enableErrorLogging = true;
296

297
        /**
298
         * 是否启用流式镜像，默认启用
299
         */
300
        private boolean enableStreamMirroring = true;
301

302
        /**
303
         * 最大镜像请求体大小（字节），默认1MB
304
         */
305
        private int maxMirrorBodySize = 1024 * 1024;
306
    }
307
}