Spring Webflux中复用初始请求TraceId到Reactor Context的正确方案

背景 #

我有一个运行在Spring Boot 3.1.0上的服务，通过REST和AMQP通信。收到REST调用后，用reactor-rabbitmq把请求体内容发布到RabbitMQ队列，之后立即返回HTTP响应。

REST控制器代码 #

@RestController
@RequiredArgsConstructor
@RequestMapping("${api.baseurl}")
public class CalculationInitiateController {

    private final RequestMapper requestMapper;
    private final ResponseMapper responseMapper;
    private final CalculationInitializer initializer;

    @ResponseStatus(code = HttpStatus.CREATED)
    @PostMapping("/initiate")
    public Mono<CalculationInitResponseDto> initiateCalculation(@RequestBody @Valid CalculationInitiationRequestDto request) {
        return Mono.just(requestMapper.map(request))
                .flatMap(initializer::initiateCalculation)
                .map(responseMapper::map);
    }
}

RabbitMQ消息发布代码 #

@Slf4j
@AllArgsConstructor
@Service
public class CalculationInitializer {

    private final Sender sender;
    private final ObjectMapper objectMapper;
    private final EventPublisherProperties eventPublisherProperties;
    private final OutboundMessageFactory messageFactory;

    public Mono<Boolean> initialize(CalculationEvent calculationEvent) {
        log.info("Publishing internal {} event for calculation: {}", eventType, calculationEvent.request().calculationId());
        var bytePayload = objectMapper.writeValueAsBytes(calculationEvent);
        var outboundMessage = new OutboundMessage("", eventConfig.getRoutingKey(), messageProperties(eventConfig), bytePayload);
        return send(outboundMessage);    
    }

    private AMQP.BasicProperties messageProperties(EventPublisherProperties.EventConfig eventConfig) {
        var context = tracer.currentTraceContext().context();
        return new AMQP.BasicProperties.Builder()
                .correlationId(context.traceId())
                .headers(Map.of(EVENT_TYPE_HEADER, eventConfig.headerValue(), "traceId", context.traceId(), "spanId", context.spanId()))
                .build();
    }

    private Mono<Boolean> send(OutboundMessage outboundMessage) {
        return sender.sendWithPublishConfirms(Mono.just(outboundMessage)).next()
                .flatMap(this::checkIfAcknowledged);
    }

    private Mono<Boolean> checkIfAcknowledged(OutboundMessageResult<OutboundMessage> result) {
        if (result.isAck()) {
            return Mono.just(Boolean.TRUE);
        } else {
            log.warn("Message not Acknowledged !!");
            return Mono.error(new IllegalStateException("Did not receive ACK on message send"));
        }
    }
}

RabbitMQ消息消费代码 #

@Slf4j
@Component
@AllArgsConstructor
public class EventMessageListener {

    private final EventMessageReceiverProperties eventMessageReceiverProperties;
    private final Receiver eventReceiver;
    private final ConsumeOptions consumeOptions;
    private final Tracer tracer;
    private final CalculationEventHandler calculationEventHandler;

    @EventListener(ApplicationReadyEvent.class)
    public void receiveMessages() {
        eventReceiver.consumeManualAck(eventMessageReceiverProperties.getQueue(), consumeOptions)
                .flatMap(this::handleMessage)
                .doFinally(s -> eventReceiver.close())
                .subscribe();
    }

    private Mono<Void> handleMessage(AcknowledgableDelivery message) {
        var traceId = message.getProperties().getHeaders().get("traceId").toString();
        var spanId = message.getProperties().getHeaders().get("spanId").toString();
        var context = tracer.traceContextBuilder().traceId(traceId).spanId(spanId).build();
        return Mono.defer(() -> calculationEventHandler.handle(message)
                .doOnSuccess(v -> message.ack())
                .onErrorResume(ex -> {
                    log.error("Failed to handle message {}", new String(message.getBody()));
                    log.error("Exception:", ex);
                    message.nack(false);
                    return Mono.empty();
                }))
                .contextWrite(Context.of(TraceContext.class, context));
    }
}

事件处理代码 #

@Component
@AllArgsConstructor
public class CalculationEventHandler  {

    private final ObjectMapper objectMapper;
    private final CalculationErrorHandler calculationErrorHandler;
    private final CalculationEventProcessor calculationEventProcessor;
    private final EventMessageReceiverProperties eventMessageReceiverProperties;

    public Mono<Void> handle(AcknowledgableDelivery message) {
        try {
            var event = objectMapper.readValue(message.getBody(), CalculationEvent.class);
            return calculationEventProcessor.process(event)
                    .onErrorResume(t -> calculationErrorHandler.handleError(message, t, event))
                    .then(Mono.empty());
        } catch (IOException e) {
            return Mono.error(e);
        }
    }
}

验证逻辑代码 #

@Service
@AllArgsConstructor
public class CalculationValidator {

    private final List<ValidationService> validationServices;

    public Mono<Void> validate(CalculationInput input) {
        return Flux.mergeDelayError(Queues.XS_BUFFER_SIZE, validateWithEachService(input).toArray(Publisher[]::new)).then();
    }

    private List<Mono<Void>> createAttributeRequests(CalculationInput input) {
        return validationServices.stream()
                .map(validationService -> validationService.validate(input))
                .toList();
    }
}

ValidationService实现示例 #

@Service
public class ThresholdValidatorServiceImpl implements ValidatorService {

    private final WebClient webClient;
    private final String endpoint;
    private final ThresholdInputMapper mapper;

    public ThresholdValidatorServiceImpl(WebClient.Builder builder,
                                        ValidationServiceErrorFilterFactory errorFilterFactory,
                                        ThresholdInputMapper mapper,
                                        @Value("${integration.threshold-validator.url}") String gateway,
                                        @Value("${integration.threshold-validator.endpoint}") String endpoint) {
        this.webClient = builder
                .baseUrl(dpGateway)
                .filter(errorFilterFactory.createFilterFor("threshold-validator"))
                .build();
        this.endpoint = endpoint;
        this.mapper = mapper;
    }

    @Override
    public Mono<JsonNode> get(CalculationInput input) {
        return Mono.just(mapper.map(input))
                .map(body -> webClient.post()
                        .uri(endpoint)
                        .content(MediaType.APPLICATION_JSON)
                        .bodyValue(body)
                        .retrieve()
                        .bodyToMono(Void.class));
    }
}

目标 #

我希望复用初始REST请求的同一个traceId——把traceId存入AMQP消息属性，Rabbit接收器消费消息后，将traceId注入Reactor Context，让每个ValidationService中通过WebClient发起的API调用都使用初始请求的traceId。

问题 #

已经成功将初始请求的traceId存入AMQP消息属性，消费消息后取出并以TraceContext形式写入Reactor Context，但后续WebClient发起API调用时，每个请求都会生成新的traceId，不符合预期。

疑问 #

这种需求能否实现？如果可以，正确的实现方式是什么？

使用的依赖 #

• io.micrometer:micrometer-tracing:1.1.2
• io.micrometer:context-propagation:1.0.3
• io.projectreactor.rabbitmq:reactor-rabbitmq:1.5.6
• org.springframework.boot:spring-boot-starter-webflux:3.1.0

解决方案 #

这种需求完全可以实现，核心是确保TraceContext正确传播到Reactor Context，并且WebClient能正确读取上下文里的Trace信息。以下是关键修正点：

1. 正确设置Reactor Context的TraceContext传播 #

当前在handleMessage方法中直接写入上下文的方式不符合Micrometer Tracing的传播规范，需要用Tracer的currentTraceContext来包装流水线，确保上下文正确传递：

private Mono<Void> handleMessage(AcknowledgableDelivery message) {
    var traceId = message.getProperties().getHeaders().get("traceId").toString();
    var spanId = message.getProperties().getHeaders().get("spanId").toString();
    var traceContext = tracer.traceContextBuilder()
            .traceId(traceId)
            .spanId(spanId)
            .build();
    
    // 使用Tracer的上下文工具包装流水线，确保TraceContext正确传播
    return tracer.currentTraceContext().wrap(Mono.defer(() -> calculationEventHandler.handle(message)
                    .doOnSuccess(v -> message.ack())
                    .onErrorResume(ex -> {
                        log.error("Failed to handle message {}", new String(message.getBody()));
                        log.error("Exception:", ex);
                        message.nack(false);
                        return Mono.empty();
                    })))
            .contextWrite(tracer.currentTraceContext().contextWrite(traceContext));
}

2. 修正WebClient调用的流水线错误 #

原ThresholdValidatorServiceImpl的get方法使用map操作创建WebClient请求，这会导致请求无法执行，且上下文无法传递。需要改为flatMap，同时修正返回类型的不匹配问题：

@Override
public Mono<Void> validate(CalculationInput input) {
    return Mono.just(mapper.map(input))
            .flatMap(body -> webClient.post()
                    .uri(endpoint)
                    .contentType(MediaType.APPLICATION_JSON)
                    .bodyValue(body)
                    .retrieve()
                    .bodyToMono(Void.class));
}

3. 确保上下文传播器配置正确 #

Spring Boot 3.x配合Micrometer Tracing 1.1.x会自动注册TraceContextPropagator，无需手动配置。如果存在自定义配置，需确保没有覆盖默认的上下文传播逻辑。

4. 验证TraceContext传播 #

可以在ValidationService中添加日志验证TraceContext是否正确读取：

@Override
public Mono<Void> validate(CalculationInput input) {
    return Mono.deferContextual(ctx -> {
        TraceContext traceContext = ctx.getOrDefault(TraceContext.class, null);
        if (traceContext != null) {
            log.info("Using traceId: {} for validation request", traceContext.traceId());
        }
        return Mono.just(mapper.map(input))
                .flatMap(body -> webClient.post()
                        .uri(endpoint)
                        .contentType(MediaType.APPLICATION_JSON)
                        .bodyValue(body)
                        .retrieve()
                        .bodyToMono(Void.class));
    });
}

关键原理 #

Micrometer Tracing通过ContextPropagator将TraceContext与Reactor Context绑定，WebClient的ObservationWebClientFilter会自动从Reactor Context中读取TraceContext，并将其作为traceparent请求头传递给下游服务，从而复用同一个traceId。

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