Creating new "child" traces in reactive process

[joachimm]

Similar to issue: 6427
I need to split traces into multiple smaller ones, due to size limits of traces in our Grafana/Tempo setup.
This time in a reactive context.
This works fine when working with Monos but not Flux:

    private  <T> Publisher<T> createNewOtelTrace(Supplier<Flux<T>> supplier) {
         return Flux.deferContextual(contextView -> {
                // Create a new span
             Tracer tracer = GlobalOpenTelemetry.getTracer("batches");
             SpanContext parentSpanContext = Span.current().getSpanContext();
             String parentTraceId = parentSpanContext.getTraceId();

             Span span = tracer.spanBuilder("download").setSpanKind(SpanKind.SERVER).setNoParent()
                 .addLink(parentSpanContext).startSpan();
             Context otelContext = Baggage.current().toBuilder().put("parent-traceid", parentTraceId).build()
                 .storeInContext(Context.current());

             try (Scope _ignore = otelContext.makeCurrent(); Scope unused = span.makeCurrent()) {
                 log.info("OTEL");
                 // Add attributes or events to the span
                    span.setAttribute("key", "value");
                    span.addEvent("Operation started");

                    // Business logic goes here
                    return supplier.get()
                        .doOnNext(item -> span.addEvent("Processing item: " + item))
                        .doOnError(span::recordException)
                        .doFinally(signal -> span.end())
                        .contextWrite(reactiveContext -> reactiveContext.put(Context.class, otelContext));
                }
            });
    }

The log: OTEL is written with the new trace/span id, but the code running in supplier is not.
I have been experimenting with defer/deferContextual and removing .contextWrite with same issue.

with Mono:

    private  <T> Publisher<T> createNewOtelTrace(Supplier<Mono<T>> supplier) {
         return Mono.deferContextual(context -> {
                // Create a new span
             Tracer tracer = GlobalOpenTelemetry.getTracer("batches");
             SpanContext parentSpanContext = Span.current().getSpanContext();
             String parentTraceId = parentSpanContext.getTraceId();

             Span span = tracer.spanBuilder("download").setSpanKind(SpanKind.SERVER).setNoParent()
                 .addLink(parentSpanContext).startSpan();
             Context otelContext = Baggage.current().toBuilder().put("parent-traceid", parentTraceId).build()
                 .storeInContext(Context.current());

             try (Scope _ignore = otelContext.makeCurrent(); Scope unused = span.makeCurrent()) {
                 log.info("OTEL");
                 // Add attributes or events to the span
                    span.setAttribute("key", "value");
                    span.addEvent("Operation started");

                    // Business logic goes here
                    return supplier.get()
                        .doOnSuccess(result -> span.addEvent("Operation completed"))
                        .doOnError(span::recordException)
                        .doFinally(signal -> span.end())
                        .contextWrite(reactiveContext -> reactiveContext.put(Context.class, otelContext));
                }
            });
    }

This works as expected with Mono, the new trace used in supplier.

How should I set up the Flux case for it to work?

[laurit]

Mono/Flux uses context that was active when it was subscribed. You could use something like

    return new Publisher<T>() {
      @Override
      public void subscribe(Subscriber<? super T> subscriber) {
        try (Scope ignore = Context.root().makeCurrent()) {
          delegate.subscribe(subscriber);
        }
      }
    };

to reset the context.
try (Scope _ignore = otelContext.makeCurrent(); Scope unused = span.makeCurrent()) {
does not make sense. No need to open a scope just to overwrite it with another one.

[joachimm]

Thanks for the reply, I think I did not adequately explain my issue.
I would like to write code like below,
Specifically I want the Supplier, given to createNewOtelTrace, using a different trace context than the rest of the chain. Due to resource limits in our Grafana/Tempo.
Currently for Mono:
Everything up until the concatMap uses one trace let's call it trace1.
consumeTile uses the new trace I created trace2
cleanUp then goes back to using trace1.
This is the behaviour I hope to see when using a Flux as well.

                Mono.just(uuid)
                .flatMapMany(this::loadUnprocessedTileInfo)
                .bufferUntilChanged(input -> Tuples.of(input.tileId(), input.inventoryId()))
                .map(TileEventHandlers::createGroup) // returns Flux<TileGroup>
                .concatMap((TileGroup group) -> createNewOtelTrace(() -> 
                // I want this to run using a new trace
                   loadTileForGroup(group) // returns Flux<Tile>
                       .flatMap(tup -> consumeTile(mapDeliveryEvent, tup))
                ))
               .collectList()
               .flatMap(this::cleanUp)

Regarding:

try (Scope _ignore = otelContext.makeCurrent(); Scope unused = span.makeCurrent()) {

I hope to activate the span I just created and somehow get the baggage I added to be part of it.
Is there a cleaner way of doing that?

[laurit]

Since mono/flux is run with the context that was used during subscribe you'll have to set your context there too. You could wrap the mono/flux 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

Line 44 in 343a7f4

public static <T> Mono<T> wrapMono(Mono<T> mono, Context context) {

Alternatively you could try

opentelemetry-java-instrumentation/instrumentation/reactor/reactor-3.1/library/src/main/java/io/opentelemetry/instrumentation/reactor/v3_1/ContextPropagationOperator.java

Line 230 in 343a7f4

public static <T> Flux<T> runWithContext(Flux<T> publisher, Context tracingContext) {

I hope to activate the span I just created and somehow get the baggage I added to be part of it.
Is there a cleaner way of doing that?

Sorry, after thinking a bit more about this I think this actually works because span.makeCurrent() uses the context from otelContext.makeCurrent(). I think it would be easier to follow if you'd use .storeInContext(Context.current().with(span)) instead storeInContext(Context.current()) then you could skip span.makeCurrent().