Add test coverage

Author: RanVaknin

aws-sdk-java/pom.xml

@@ -2113,6 +2113,11 @@ Amazon AutoScaling, etc).</description>
             <artifactId>bcmdashboards</artifactId>
             <version>${awsjavasdk.version}</version>
         </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>utils-lite</artifactId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
     </dependencies>
     <build>
         <finalName>${project.artifactId}-${project.version}</finalName>

test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java

@@ -18,6 +18,10 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider;
 import software.amazon.awssdk.awscore.interceptor.TraceIdExecutionInterceptor;
@@ -243,7 +247,7 @@ public void beforeMarshalling(Context.BeforeMarshalling context, ExecutionAttrib
     }
 
     @Test
-    public void traceIdInterceptorWithNewThreadInheritsTraceId() throws Exception {
+    public void traceIdInterceptorWithNewThreadInheritsTraceId() {
         EnvironmentVariableHelper.run(env -> {
             env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
 
@@ -275,4 +279,67 @@ public void traceIdInterceptorWithNewThreadInheritsTraceId() throws Exception {
             }
         });
     }
+
+    @Test
+    public void traceIdInterceptorWithExecutiveServicePreservesTraceId() {
+        EnvironmentVariableHelper.run(env -> {
+            env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
+
+            SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123");
+            ExecutorService executor = Executors.newFixedThreadPool(2);
+            try (MockSyncHttpClient mockHttpClient = new MockSyncHttpClient();
+                 ProtocolRestJsonClient client = ProtocolRestJsonClient.builder()
+                                                                       .region(Region.US_WEST_2)
+                                                                       .credentialsProvider(AnonymousCredentialsProvider.create())
+                                                                       .httpClient(mockHttpClient)
+                                                                       .build()) {
+
+                mockHttpClient.stubNextResponse(HttpExecuteResponse.builder()
+                                                                   .response(SdkHttpResponse.builder().statusCode(200).build())
+                                                                   .responseBody(AbortableInputStream.create(new StringInputStream("{}")))
+                                                                   .build());
+
+                executor.submit(() -> client.allTypes()).get();
+
+                List<SdkHttpRequest> requests = mockHttpClient.getRequests();
+                assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123");
+
+            } catch (InterruptedException | ExecutionException e) {
+                throw new RuntimeException(e);
+            } finally {
+                SdkInternalThreadLocal.clear();
+            }
+        });
+    }
+
+    @Test
+    public void traceIdInterceptorWithRunAsyncDoesNotPreservesTraceId() throws Exception {
+        EnvironmentVariableHelper.run(env -> {
+            env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
+
+            SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123");
+            try (MockSyncHttpClient mockHttpClient = new MockSyncHttpClient();
+                 ProtocolRestJsonClient client = ProtocolRestJsonClient.builder()
+                                                                       .region(Region.US_WEST_2)
+                                                                       .credentialsProvider(AnonymousCredentialsProvider.create())
+                                                                       .httpClient(mockHttpClient)
+                                                                       .build()) {
+
+                mockHttpClient.stubNextResponse(HttpExecuteResponse.builder()
+                                                                   .response(SdkHttpResponse.builder().statusCode(200).build())
+                                                                   .responseBody(AbortableInputStream.create(new StringInputStream("{}")))
+                                                                   .build());
+
+                CompletableFuture.runAsync(client::allTypes).get();
+
+                List<SdkHttpRequest> requests = mockHttpClient.getRequests();
+                assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).isEmpty();
+
+            } catch (InterruptedException | ExecutionException e) {
+                throw new RuntimeException(e);
+            } finally {
+                SdkInternalThreadLocal.clear();
+            }
+        });
+    }
 }