Issue with Trace Context Sharing Across HTTP Requests in Spring WebFlux Application

[srinusanduri]

Hi,

We've encountered an unusual issue with the instrumentation of a Spring WebFlux application, where a single trace context is being shared across multiple HTTP requests. This issue arises when we add manual instrumentation to set span attributes and baggage. Notably, we do not manually create spans. If we remove this manual instrumentation, the issue disappears.

Here is an example of a trace without manual instrumentation, which behaves as expected, containing eight spans:

Conversely, here is a trace with manual instrumentation. This trace incorrectly grows in size, accumulating thousands of spans as more HTTP requests are made to the Spring WebFlux application. Despite expectations of only eight spans, it appears that all child spans from the requests are being added to a single trace, indicating an issue with trace separation:

Below is the snippet of code with manual instrumentation:

    public Mono<CommandStatus> sendCommand(
            CommandServiceRequest commandServiceRequest, String userName, String clientId, String accountId) {
        return Mono
                .just(SendVehicleCommandContext.fromRequest(userName, accountId, clientId, commandServiceRequest))
                .map(context -> context.setSupportedCommandType(COMMAND_TYPE))
                .map(context -> {
                    var commandPermission = supportedCommandTypePermissionProvider.getCommandPermission(context.getSupportedCommandType());
                    var commandId = randomSupport.randomUUID();

                    // Set span attribute and baggage
                    Baggage.current().toBuilder().put("command_id", String.valueOf(commandId)).build().makeCurrent();
                    Span.current().setAttribute("command_id", String.valueOf(commandId));

                    return context.setCommandId(commandId).setCommandPermission(commandPermission);
                })
                .doOnNext(context -> log.debug("send command {} for account {}", context.getCommandId(), context.getAccountId()))
                .flatMap(context -> assetsService.getAssetForAccount(context.getAccountId())
                        .doOnNext(asset -> {

                            // Set baggage
                            Baggage.current().toBuilder().put("asset_id", asset.getId().toString()).build().makeCurrent();

                        })
                        .map(context::setVehicleAsset))
                .flatMap(context -> iamService
                        .validateCommandAuthorization(context.getUsername(), context.getAccountId(),
                                context.getCommandPermission().getIssuePermission())
                        .thenReturn(context))
                .flatMap(context -> {
                    var processedCommandServiceRequest = commandServiceRequestProcessor.processCommandServiceRequest(
                            context.getCommandServiceRequest(),
                            context.getCommandId(), userName);

                    String commandId = String.valueOf(context.getCommandId());

                    // Set span attribute and baggage
                    Baggage.current().toBuilder().put("command_id", commandId).build().makeCurrent();
                    Span.current().setAttribute("command_id", commandId);

                    return issueCommand(processedCommandServiceRequest, context.getClientId(), context.getAccountId())
                            .map(context::setCommandStatus);
                })
                .map(SendCommandContext::toResponse)
                .doOnNext(commandStatus -> {
                    log.debug("sent command {} for account {}", commandStatus.getId(), accountId);
                });
    }

Could this unusual behavior be related to the manual instrumentation code provided? We would greatly appreciate any insights or suggestions you might have.

Thank you for your time and assistance.

[laurit]

It is a bug in your manual instrumentation. Scope returned from makeCurrent() must be closed. Failure to do so may result in the observed behavior.

[srinusanduri]

Thank you, @laurit.

Could you advise on how to safely close or alter the context in the code snippet above? Given that the current span/context is a root span created automatically by the Netty library's instrumentation, I'm concerned about potentially disrupting the context for the remainder of the execution flow. Is there a way to modify the context without negatively impacting the rest of the application's trace continuity?

[laurit]

I would try wrapping the mono with something like

opentelemetry-java-instrumentation/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/AdviceUtils.java

Lines 56 to 72 in b93537d

	private static class ContextMono<T> extends Mono<T> {
	private final Mono<T> delegate;
	private final Context parentContext;
	ContextMono(Mono<T> delegate, Context parentContext) {
	this.delegate = delegate;
	this.parentContext = parentContext;
	}
	@Override
	public void subscribe(CoreSubscriber<? super T> actual) {
	try (Scope ignored = parentContext.makeCurrent()) {
	delegate.subscribe(actual);
	}
	}
	}

The parent context would be current context + the baggage you wish to add Context context = Baggage.current().toBuilder().put("command_id", String.valueOf(commandId)).build().storeInContext(Context.current()); Context from the time when mono is subscribed is propagated automatically.

[srinusanduri]

Thank you for the suggestions, @laurit.

I have updated the code snippet with ContextMono to manage the trace context. However, I'd appreciate it if you could review the updated code to verify if the ContextMono usage for setting the baggage is implemented correctly.

    public Mono<CommandStatus> sendCommand(
            CommandServiceRequest commandServiceRequest, String userName, String clientId, String accountId) {
        return Mono
                .just(SendVehicleCommandContext.fromRequest(userName, accountId, clientId, commandServiceRequest))
                .map(context -> context.setSupportedCommandType(COMMAND_TYPE))
                .map(context -> {
                    var commandPermission = supportedCommandTypePermissionProvider.getCommandPermission(context.getSupportedCommandType());
                    var commandId = randomSupport.randomUUID();

                    SendVehicleCommandContext sendCommandContext = context.setCommandId(commandId).setCommandPermission(commandPermission);

                    // Set span attribute and baggage
                    Context otelContext = Baggage.current().toBuilder()
                            .put("command_id", String.valueOf(commandId))
                            .build()
                            .storeInContext(Context.current());

                    Span.current().setAttribute("command_id", String.valueOf(commandId));

                    return new ContextMono<>(Mono.just(sendCommandContext), otelContext);
                })
                .doOnNext(context -> log.debug("send command {} for account {}", context.getCommandId(), context.getAccountId()))
                .flatMap(context -> assetsService.getAssetForAccount(context.getAccountId())
                        .map(asset -> {

                            // Set span attribute and baggage
                            Context otelContext = Baggage.current().toBuilder()
                                    .put(Keys.ASSET_ID, asset.getId().toString())
                                    .build()
                                    .storeInContext(Context.current());
                            Span.current().setAttribute(Keys.ASSET_ID, asset.getId().toString());

                            return new ContextMono<>(Mono.just(context.setVehicleAsset(asset)), otelContext);
                        })
                        .flatMap(contextMono -> contextMono)
                )
                .flatMap(context -> iamService
                        .validateCommandAuthorization(context.getUsername(), context.getAccountId(),
                                context.getCommandPermission().getIssuePermission())
                        .thenReturn(context))
                .flatMap(context -> {
                    var processedCommandServiceRequest = commandServiceRequestProcessor.processCommandServiceRequest(
                            context.getCommandServiceRequest(), context.getCommandId(), userName);
                    return issueCommand(processedCommandServiceRequest, context.getClientId(), context.getAccountId())
                            .map(context::setCommandStatus);
                })
                .map(SendCommandContext::toResponse)
                .doOnNext(commandStatus -> {
                    log.debug("sent command {} for account {}", commandStatus.getId(), accountId);
                });
    }

Also, I would like to confirm my understanding regarding the scope closure in the ContextMono.subscribe() method. Based on the implementation, it appears that the context is correctly closed after its use within the try-with-resources block. Additionally, I assume that the parent span context remains active for the remainder of the flow beyond this block. Is this understanding correct? Here is the relevant code snippet:

@Override 
public void subscribe(CoreSubscriber<? super T> actual) { 
    try (Scope ignored = parentContext.makeCurrent()) { 
        delegate.subscribe(actual); 
    } 
}

I find this approach to setting baggage a bit verbose and complex, which may pose challenges for teams, especially when integrating it into our existing codebase. Could you suggest any more convenient utility methods or best practices for setting baggage more efficiently?

[laurit]

Additionally, I assume that the parent span context remains active for the remainder of the flow beyond this block.

parentContext is active only inside the try block, after the try block current context reverts to what it was before. Instrumentation runs mono with the context that was active during the subscribe so I guess you could say that parent span context remains active for the remainder of the flow beyond this block.

I find this approach to setting baggage a bit verbose and complex, which may pose challenges for teams, especially when integrating it into our existing codebase. Could you suggest any more convenient utility methods or best practices for setting baggage more efficiently?

I really know next to nothing about webflux and reactor, I'd assume that someone with good knowledge of these frameworks can come up with appropriate utilities to make this easier.