SageMakerRuntimeAsyncClient invokation generates a new root trace and not part of existing trace.

[chenlin2001]

I have a Spring application (3.5.7) which is instrumented with the opentelemetry java agent 2.18.1
I have the following Bean:

 @Bean
    public SageMakerRuntimeAsyncClient sageMakerRuntimeAsyncClient() {
        return SageMakerRuntimeAsyncClient.builder()
                .credentialsProvider(DefaultCredentialsProvider.create())
                .region(Region.of(awsRegion))
                .asyncConfiguration(b -> b.advancedOption(
                        SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR,
                        Executors.newVirtualThreadPerTaskExecutor()))
                .build();
    }

Call come from controller -> sqs producer -> sqs consumer -> Sagemaker
In GarafanCloud I can see 2 traces - one with 3 spans controller -> sps producer -> sqs consumer. And one for SageMaker (SageMakerRuntime.InvokeEndpoint)

What is the recommended way using AsyncClient + virtual threads to get one trace and not the SageMakerRuntime.InvokeEndpoint as its root trace?

I tried the following also:

1. Wrapping the virtual threads with delagator which has the current context
2. Using NettyNioAsyncHttpClient for the http cliennt:

Neither solved the issue

    @Bean(destroyMethod = "close")
    public ExecutorService vthreadExecutor() {
        return Executors.newVirtualThreadPerTaskExecutor();
    }

    // 2. Wrap executor only when OTEL is enabled
    @Bean
    public Executor contextAwareExecutor(ExecutorService vthreadExecutor) {
        if (!openTelemetryEnabled) {
            return vthreadExecutor;
        }

        return task -> {
            Context parent = Context.current();
            vthreadExecutor.execute(() -> {
                try (Scope s = parent.makeCurrent()) {
                    task.run();
                }
            });
        };
    }

    // 3. Use context-aware executor for SageMaker async callbacks
    @Bean
    public SageMakerRuntimeAsyncClient sageMakerRuntimeAsyncClient(
            Executor contextAwareExecutor) {

        return SageMakerRuntimeAsyncClient.builder()
                .httpClientBuilder(NettyNioAsyncHttpClient.builder())
                .credentialsProvider(DefaultCredentialsProvider.create())
                .region(Region.of(awsRegion))
                .asyncConfiguration(b -> b.advancedOption(
                        SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR,
                        contextAwareExecutor))
                .build();
    }

[laurit]

We don't have any tests for SageMakerRuntimeAsyncClient. It could be that the issue is specific to SageMakerRuntimeAsyncClient or it could be that it has something to do with using httpClientBuilder(NettyNioAsyncHttpClient.builder()) or .asyncConfiguration(b -> b.advancedOption(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR, contextAwareExecutor)). We should figure out what exactly causes the issue and add tests for it.