Proper way of passing manually generated Context when using CompletableFuture

[zbytt]

Hello !
In my current setup, I'm using instrumentation agent with XRay propagator (aws-distro) + xray exporter in collector.

Application fetches messages from SQS (in batches), every message is then processed one by one asynchronously.
Default trace created by agent in this situation is the one taken from X-Amzn-Trace-Id response from SQS, which is the same for all messages in that batch.

As every message in SQS contains AwsTraceHeader attribute, I tried to use it to create new context :

    final String traceFromMessage = getMsgTraceId(message);
    final Context context = GlobalOpenTelemetry.getPropagators()
            .getTextMapPropagator()
            .extract(Context.current(),
                    //put trace from message under X-Amzn-Trace-Id, so that XrayPropagator can pick it up
                    Map.of("X-Amzn-Trace-Id", traceFromMessage),
                    new TextMapGetter<>() {
                        @Override
                        public Iterable<String> keys(final Map<String, String> carrier) {
                            return carrier.keySet();
                        }

                        @Override
                        public String get(final Map<String, String> carrier, final String key) {
                            return carrier.get(key);
                        }
                    });`

In that setup created context must be marked as current, and then closed somewhere.

Next steps are executed using CompletableFuture/Executors framework, which makes it a bit more complicated.
How should I approach the execution of makeCurrent()/close() methods in that case ?

Given very simple example of :

 CompletableFuture.supplyAsync(() -> {
                //do something with the msg
                return processMessage(message);
            }, executor)
            .thenAccept(msg -> {
                ack(msg);
            })
            .whenComplete((msg, exc) -> {
                if (exc != null) {
                    LOG.info("Processing failed.");
                } else {
                    LOG.info("Processing succeeded.");
                }
            });

Should I pass the context around manually as an instance and call makeCurrent()/close() in places where it is needed like :

    CompletableFuture.supplyAsync(() -> {
                //do something with the msg
                try (final Scope ignore = context.makeCurrent()) {
                    return processMessage(message);
                }
            }, executor)
            .thenAccept(msg -> {
                try (final Scope ignore = context.makeCurrent()) {
                    ack(msg);
                }
            })
            .whenComplete((msg, exc) -> {
                try (final Scope ignore = context.makeCurrent()) {
                    if (exc != null) {
                        LOG.info("Processing failed.");
                    } else {
                        LOG.info("Processing succeeded.");
                    }
                }
            });

Is this the right (but not very convenient) way to do it? Will this approach create a proper spans ?
(Idea based on the comment here : #7840 (comment) )

(I've tried to use this one, looking at the logs - it looks quite right - all the logs have proper trace_id in MDC, but this 'custom' trace is missing in service map in xray (weirdly though, the only missing part is the one from application that was creating the custom span, all the calls to external apps made from that app are in place, with proper 'custom' trace propagated )

Is there a better/more convenient way to achieve this ?

[laurit]

Activating the scope the way you are doing it is definitely the most reliable way. Java agent also has automatic propagation for completable futures where callbacks like thenAccept are called with the context where the future was completed. Note that if thenAccept is called on an already completed future then thenAccept will be called with the current context. So the alternative you could try is

try (final Scope ignore = context.makeCurrent()) {
        CompletableFuture.supplyAsync(() -> {
                //do something with the msg
                return processMessage(message);
            }, executor)
            .thenAccept(msg -> {
                ack(msg);
            })
            .whenComplete((msg, exc) -> {
                if (exc != null) {
                    LOG.info("Processing failed.");
                } else {
                    LOG.info("Processing succeeded.");
                }
        });
}

and see whether it works for you.

[zbytt]

Thanks a lot for your input !
I tried the single try with resources pattern you proposed even before creating this discussion, but it did not work for me as I would expect - traces created this way were crashing AWS Xray traces map. I assumed that I just used it wrong, hence the question here.
If you think this should work, I will go back to it and try to figure out what is causing the crash.

Note that if thenAccept is called on an already completed future then thenAccept will be called with the current context.

I will keep that in mind, but I do not think this can break anything in my case, thanks for bringing that up.

Thanks again.