Spring WebFlux中正确的等待实现方案咨询

问题背景 #

我们有一个HTTP端点，接收命令后通过Kafka发送至其他服务，必须等待Kafka返回响应后，再向客户端返回ResponseEntity。现有代码使用CountDownLatch阻塞等待Kafka响应，违反了WebFlux非阻塞的设计原则，会占用IO线程导致无法处理其他请求，但客户端流程无法修改，需要提供符合WebFlux规范的实现方式。

现有代码的核心问题 #

现有waitForAuthentication方法是阻塞式的，在WebFlux的反应式链中调用会直接占用Netty事件循环线程，导致线程池耗尽，严重影响应用吞吐量和响应性能。

解决方案：基于Reactor的非阻塞改造 #

核心思路是将阻塞的等待逻辑包装为Reactor的Mono类型，将阻塞操作转移到专门的弹性线程池，同时通过请求ID关联Kafka请求与响应，确保消息匹配。

1. 改造认证服务：返回非阻塞Mono #

将原阻塞方法改为返回Mono<String>，使用Mono.create包装阻塞逻辑，并通过subscribeOn指定弹性线程池：

public Mono<String> waitForAuthentication(String data) {
    // 生成唯一请求ID，用于匹配Kafka响应
    String requestId = UUID.randomUUID().toString();
    
    return Mono.create(sink -> {
        CountDownLatch countDownLatch = new CountDownLatch(1);
        Listener listener = new Listener(countDownLatch, requestId, sink);
        
        // 订阅Kafka响应主题
        subscribe("topic_abc", String.class, listener);
        
        // 发送携带requestId的认证请求到Kafka
        sendAuthenticationRequestToKafka(data, requestId);
        
        try {
            boolean isSuccess = countDownLatch.await(2, TimeUnit.SECONDS);
            if (!isSuccess) {
                sink.error(new CustomException("Authentication timeout"));
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            sink.error(new CustomException("Authentication interrupted"));
        } catch (Exception e) {
            sink.error(new CustomException("Authentication failed"));
        }
    }).subscribeOn(Schedulers.boundedElastic()); // 将阻塞操作转移到弹性线程池
}

2. 改造监听器：按请求ID匹配响应 #

修改监听器，增加请求ID和MonoSink，仅处理与当前请求ID匹配的Kafka响应：

@RequiredArgsConstructor
public class Listener implements Consumer<String> {
    private final CountDownLatch countDownLatch;
    private final String requestId;
    private final MonoSink<String> sink;

    @Override
    public void accept(String responseData) {
        // 从响应数据中解析出请求ID（需根据实际响应格式实现）
        String responseRequestId = parseRequestId(responseData);
        if (requestId.equals(responseRequestId)) {
            sink.success(responseData);
            countDownLatch.countDown();
            // 取消当前监听器订阅，避免接收无关消息
            unsubscribe(this);
        }
    }
    
    // 示例：解析响应中的requestId（根据实际格式调整）
    private String parseRequestId(String responseData) {
        // 假设响应是JSON格式，比如{"requestId":"xxx", "data":"xxx"}
        ObjectMapper objectMapper = new ObjectMapper();
        try {
            JsonNode node = objectMapper.readTree(responseData);
            return node.get("requestId").asText();
        } catch (IOException e) {
            return "";
        }
    }
}

3. 改造服务层：使用反应式链处理结果 #

原服务层的阻塞调用改为订阅Mono，通过反应式操作符处理结果和异常：

public Mono<String> requestProcess(byte[] b) {
    return authenticationService.waitForAuthentication(parseBytes(b))
            .filter(StringUtils::isNotEmpty)
            .switchIfEmpty(Mono.error(new CustomException("Failed to authenticate")));
}

4. 控制器保持不变 #

控制器本身已经返回Mono<ResponseEntity<byte[]>>，无需修改：

public Mono<ResponseEntity<byte[]>> post(@RequestBody Mono<byte[]> b) {
    return b.flatMap(
            identifyType -> service.requestProcess(identifyType).flatMap(this::toResponse));
}

关键设计说明 #

• 非阻塞线程模型：通过subscribeOn(Schedulers.boundedElastic())将阻塞的await操作转移到弹性线程池，避免占用WebFlux的IO线程（Netty事件循环线程），保证其他请求能正常处理。
• 请求响应匹配：通过唯一请求ID关联Kafka请求和响应，避免不同请求的消息交叉干扰。
• 异常与超时处理：通过MonoSink传递成功/失败信号，统一处理超时、中断等异常场景。

注意事项 #

• 确保Kafka订阅支持动态添加/移除监听器，避免内存泄漏。
• 若Kafka响应量较大，可结合消费者分组或分区策略优化消息匹配效率。
• 弹性线程池的大小需根据业务负载调整，避免线程过多导致资源耗尽。

内容的提问来源于stack exchange，提问作者SoT