From cbb1dd440c5bac4ad9bbe8757f7955213d0fa077 Mon Sep 17 00:00:00 2001 From: Ran Vaknin <50976344+RanVaknin@users.noreply.github.com> Date: Fri, 5 Sep 2025 09:41:57 -0700 Subject: [PATCH 01/20] Add utils-lite package (#6395) Add utils-lite package --- .brazil.json | 1 + bom/pom.xml | 5 ++ pom.xml | 2 + test/architecture-tests/pom.xml | 5 ++ .../archtests/UtilsLitePackageTest.java | 43 ++++++++++ test/tests-coverage-reporting/pom.xml | 5 ++ utils-lite/pom.xml | 79 +++++++++++++++++++ .../utilslite/SdkInternalThreadLocal.java | 55 +++++++++++++ .../awssdk/utilslite/ThreadStorageTest.java | 70 ++++++++++++++++ 9 files changed, 265 insertions(+) create mode 100644 test/architecture-tests/src/test/java/software/amazon/awssdk/archtests/UtilsLitePackageTest.java create mode 100644 utils-lite/pom.xml create mode 100644 utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java create mode 100644 utils-lite/src/test/java/software/amazon/awssdk/utilslite/ThreadStorageTest.java diff --git a/.brazil.json b/.brazil.json index 7d240cb9bf5c..362888889178 100644 --- a/.brazil.json +++ b/.brazil.json @@ -32,6 +32,7 @@ "s3-transfer-manager": { "packageName": "AwsJavaSdk-S3-TransferManager" }, "s3-event-notifications": { "packageName": "AwsJavaSdk-S3-EventNotifications" }, "sdk-core": { "packageName": "AwsJavaSdk-Core" }, + "utils-lite": { "packageName": "AwsJavaSdk-UtilsLite" }, "url-connection-client": { "packageName": "AwsJavaSdk-HttpClient-UrlConnectionClient" }, "utils": { "packageName": "AwsJavaSdk-Core-Utils" }, "imds": { "packageName": "AwsJavaSdk-Imds" }, diff --git a/bom/pom.xml b/bom/pom.xml index 18f4e532fb89..2976999b5add 100644 --- a/bom/pom.xml +++ b/bom/pom.xml @@ -212,6 +212,11 @@ utils ${awsjavasdk.version} + + software.amazon.awssdk + utils-lite + ${awsjavasdk.version} + software.amazon.awssdk cloudwatch-metric-publisher diff --git a/pom.xml b/pom.xml index 0829651b50fc..86e4f37b0ff4 100644 --- a/pom.xml +++ b/pom.xml @@ -67,6 +67,7 @@ metric-publishers release-scripts utils + utils-lite codegen-lite codegen-lite-maven-plugin archetypes @@ -664,6 +665,7 @@ cloudwatch-metric-publisher emf-metric-logging-publisher utils + utils-lite imds retries retries-spi diff --git a/test/architecture-tests/pom.xml b/test/architecture-tests/pom.xml index 4d4483104df0..9e30c195ac6c 100644 --- a/test/architecture-tests/pom.xml +++ b/test/architecture-tests/pom.xml @@ -61,6 +61,11 @@ software.amazon.awssdk ${awsjavasdk.version} + + utils-lite + software.amazon.awssdk + ${awsjavasdk.version} + s3 software.amazon.awssdk diff --git a/test/architecture-tests/src/test/java/software/amazon/awssdk/archtests/UtilsLitePackageTest.java b/test/architecture-tests/src/test/java/software/amazon/awssdk/archtests/UtilsLitePackageTest.java new file mode 100644 index 000000000000..437094a4c58f --- /dev/null +++ b/test/architecture-tests/src/test/java/software/amazon/awssdk/archtests/UtilsLitePackageTest.java @@ -0,0 +1,43 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.archtests; + +import static com.tngtech.archunit.lang.syntax.ArchRuleDefinition.classes; + +import com.tngtech.archunit.core.domain.JavaClasses; +import com.tngtech.archunit.core.importer.ClassFileImporter; +import com.tngtech.archunit.lang.ArchRule; +import org.junit.jupiter.api.Test; + +/** + * Architecture tests for the utils-lite package to ensure it only contains allowed classes. + */ +public class UtilsLitePackageTest { + + private static final JavaClasses CLASSES = new ClassFileImporter() + .importPackages("software.amazon.awssdk.utilslite"); + + @Test + public void utilsLitePackage_shouldOnlyContainAllowedClasses() { + ArchRule rule = classes() + .that().resideInAPackage("software.amazon.awssdk.utilslite") + .should().haveNameMatching(".*\\.(SdkInternalThreadLocal|SdkInternalThreadLocalTest)") + .allowEmptyShould(true) + .because("utils-lite package should only contain SdkInternalThreadLocal and its test"); + + rule.check(CLASSES); + } +} \ No newline at end of file diff --git a/test/tests-coverage-reporting/pom.xml b/test/tests-coverage-reporting/pom.xml index 891b0bb3f1ff..c3af80154ff5 100644 --- a/test/tests-coverage-reporting/pom.xml +++ b/test/tests-coverage-reporting/pom.xml @@ -326,6 +326,11 @@ http-client-spi ${awsjavasdk.version} + + software.amazon.awssdk + utils-lite + ${awsjavasdk.version} + diff --git a/utils-lite/pom.xml b/utils-lite/pom.xml new file mode 100644 index 000000000000..9642803ccd58 --- /dev/null +++ b/utils-lite/pom.xml @@ -0,0 +1,79 @@ + + + + + 4.0.0 + + software.amazon.awssdk + aws-sdk-java-pom + 2.33.3-SNAPSHOT + + utils-lite + AWS Java SDK :: Utils Lite + + A package providing minimal external utils. + + https://aws.amazon.com/sdkforjava + + + + + software.amazon.awssdk + bom-internal + ${project.version} + pom + import + + + + + + + software.amazon.awssdk + annotations + ${awsjavasdk.version} + + + org.junit.jupiter + junit-jupiter + test + + + org.assertj + assertj-core + test + + + + + + + org.apache.maven.plugins + maven-jar-plugin + + + + software.amazon.awssdk.utilslite + + + + + + + + \ No newline at end of file diff --git a/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java new file mode 100644 index 000000000000..34f3b5c9df30 --- /dev/null +++ b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java @@ -0,0 +1,55 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.utilslite; + +import java.util.HashMap; +import java.util.Map; +import software.amazon.awssdk.annotations.SdkProtectedApi; + +/** + * Utility for thread-local context storage. + */ +@SdkProtectedApi +public final class SdkInternalThreadLocal { + private static final ThreadLocal> STORAGE = ThreadLocal.withInitial(HashMap::new); + + private SdkInternalThreadLocal() { + } + + public static void put(String key, String value) { + if (value == null) { + STORAGE.get().remove(key); + } else { + STORAGE.get().put(key, value); + } + } + + public static String get(String key) { + return STORAGE.get().get(key); + } + + public static String remove(String key) { + return STORAGE.get().remove(key); + } + + public static void clear() { + STORAGE.get().clear(); + } + + public static boolean containsKey(String key) { + return STORAGE.get().containsKey(key); + } +} \ No newline at end of file diff --git a/utils-lite/src/test/java/software/amazon/awssdk/utilslite/ThreadStorageTest.java b/utils-lite/src/test/java/software/amazon/awssdk/utilslite/ThreadStorageTest.java new file mode 100644 index 000000000000..07e5e34e2ecc --- /dev/null +++ b/utils-lite/src/test/java/software/amazon/awssdk/utilslite/ThreadStorageTest.java @@ -0,0 +1,70 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.utilslite; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +class SdkInternalThreadLocalTest { + + @AfterEach + void cleanup() { + SdkInternalThreadLocal.clear(); + } + + @Test + void putAndGet_shouldStoreAndRetrieveValue() { + SdkInternalThreadLocal.put("test-key", "test-value"); + + assertThat(SdkInternalThreadLocal.get("test-key")).isEqualTo("test-value"); + } + + @Test + void get_withNonExistentKey_shouldReturnNull() { + assertThat(SdkInternalThreadLocal.get("non-existent")).isNull(); + } + + @Test + void put_withValidKeyValue_shouldStoreValue() { + SdkInternalThreadLocal.put("test-key", "test-value"); + + String removed = SdkInternalThreadLocal.remove("test-key"); + + assertThat(removed).isEqualTo("test-value"); + assertThat(SdkInternalThreadLocal.get("test-key")).isNull(); + } + + @Test + void remove_withExistingKey_shouldRemoveAndReturnValue() { + SdkInternalThreadLocal.put("test-key", "test-value"); + SdkInternalThreadLocal.put("test-key", null); + + assertThat(SdkInternalThreadLocal.get("test-key")).isNull(); + } + + @Test + void clear_withMultipleValues_shouldRemoveAllValues() { + SdkInternalThreadLocal.put("key1", "value1"); + SdkInternalThreadLocal.put("key2", "value2"); + + SdkInternalThreadLocal.clear(); + + assertThat(SdkInternalThreadLocal.get("key1")).isNull(); + assertThat(SdkInternalThreadLocal.get("key2")).isNull(); + } +} \ No newline at end of file From 3a5859eaa7b2e60cc204a71556c7325297f7c97f Mon Sep 17 00:00:00 2001 From: Ran Vaknin <50976344+RanVaknin@users.noreply.github.com> Date: Fri, 5 Sep 2025 10:57:51 -0700 Subject: [PATCH 02/20] Add support for concurrent trace id propagation (#6403) Add support for concurrent trace id propagation --- core/aws-core/pom.xml | 5 + .../TraceIdExecutionInterceptor.java | 40 +++- .../TraceIdExecutionInterceptorTest.java | 73 +++++++ .../amazon/awssdk/services/TraceIdTest.java | 187 +++++++++++++++++- 4 files changed, 300 insertions(+), 5 deletions(-) diff --git a/core/aws-core/pom.xml b/core/aws-core/pom.xml index 3135accad1a4..694098d54227 100644 --- a/core/aws-core/pom.xml +++ b/core/aws-core/pom.xml @@ -113,6 +113,11 @@ software.amazon.eventstream eventstream + + software.amazon.awssdk + utils-lite + ${awsjavasdk.version} + software.amazon.awssdk diff --git a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java index 95224228cfb4..d2b422c940eb 100644 --- a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java +++ b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java @@ -19,10 +19,12 @@ import software.amazon.awssdk.annotations.SdkProtectedApi; import software.amazon.awssdk.awscore.internal.interceptor.TracingSystemSetting; import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttribute; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; import software.amazon.awssdk.http.SdkHttpRequest; import software.amazon.awssdk.utils.SystemSetting; +import software.amazon.awssdk.utilslite.SdkInternalThreadLocal; /** * The {@code TraceIdExecutionInterceptor} copies the trace details to the {@link #TRACE_ID_HEADER} header, assuming we seem to @@ -32,27 +34,57 @@ public class TraceIdExecutionInterceptor implements ExecutionInterceptor { private static final String TRACE_ID_HEADER = "X-Amzn-Trace-Id"; private static final String LAMBDA_FUNCTION_NAME_ENVIRONMENT_VARIABLE = "AWS_LAMBDA_FUNCTION_NAME"; + private static final String CONCURRENT_TRACE_ID_KEY = "AWS_LAMBDA_X_TRACE_ID"; + private static final ExecutionAttribute TRACE_ID = new ExecutionAttribute<>("TraceId"); + + @Override + public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) { + String traceId = SdkInternalThreadLocal.get(CONCURRENT_TRACE_ID_KEY); + if (traceId != null) { + executionAttributes.putAttribute(TRACE_ID, traceId); + } + } @Override public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { Optional traceIdHeader = traceIdHeader(context); if (!traceIdHeader.isPresent()) { Optional lambdafunctionName = lambdaFunctionNameEnvironmentVariable(); - Optional traceId = traceId(); + Optional traceId = traceId(executionAttributes); if (lambdafunctionName.isPresent() && traceId.isPresent()) { return context.httpRequest().copy(r -> r.putHeader(TRACE_ID_HEADER, traceId.get())); } } - return context.httpRequest(); } + @Override + public void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) { + saveTraceId(executionAttributes); + } + + @Override + public void onExecutionFailure(Context.FailedExecution context, ExecutionAttributes executionAttributes) { + saveTraceId(executionAttributes); + } + + private static void saveTraceId(ExecutionAttributes executionAttributes) { + String traceId = executionAttributes.getAttribute(TRACE_ID); + if (traceId != null) { + SdkInternalThreadLocal.put(CONCURRENT_TRACE_ID_KEY, executionAttributes.getAttribute(TRACE_ID)); + } + } + private Optional traceIdHeader(Context.ModifyHttpRequest context) { return context.httpRequest().firstMatchingHeader(TRACE_ID_HEADER); } - private Optional traceId() { + private Optional traceId(ExecutionAttributes executionAttributes) { + Optional traceId = Optional.ofNullable(executionAttributes.getAttribute(TRACE_ID)); + if (traceId.isPresent()) { + return traceId; + } return TracingSystemSetting._X_AMZN_TRACE_ID.getStringValue(); } @@ -61,4 +93,4 @@ private Optional lambdaFunctionNameEnvironmentVariable() { return SystemSetting.getStringValueFromEnvironmentVariable(LAMBDA_FUNCTION_NAME_ENVIRONMENT_VARIABLE); // CHECKSTYLE:ON } -} +} \ No newline at end of file diff --git a/core/aws-core/src/test/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptorTest.java b/core/aws-core/src/test/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptorTest.java index b3f965a490fc..3c18d064cd0d 100644 --- a/core/aws-core/src/test/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptorTest.java +++ b/core/aws-core/src/test/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptorTest.java @@ -28,6 +28,7 @@ import software.amazon.awssdk.http.SdkHttpMethod; import software.amazon.awssdk.http.SdkHttpRequest; import software.amazon.awssdk.testutils.EnvironmentVariableHelper; +import software.amazon.awssdk.utilslite.SdkInternalThreadLocal; public class TraceIdExecutionInterceptorTest { @Test @@ -111,6 +112,78 @@ public void headerNotAddedIfNoTraceIdEnvVar() { }); } + @Test + public void modifyHttpRequest_whenMultiConcurrencyModeWithInternalThreadLocal_shouldAddTraceIdHeader() { + EnvironmentVariableHelper.run(env -> { + resetRelevantEnvVars(env); + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123"); + + try { + TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor(); + ExecutionAttributes executionAttributes = new ExecutionAttributes(); + + interceptor.beforeExecution(null, executionAttributes); + Context.ModifyHttpRequest context = context(); + + SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes); + assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + } finally { + SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID"); + } + }); + } + + @Test + public void modifyHttpRequest_whenMultiConcurrencyModeWithBothInternalThreadLocalAndSystemProperty_shouldUseInternalThreadLocalValue() { + EnvironmentVariableHelper.run(env -> { + resetRelevantEnvVars(env); + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + + SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123"); + Properties props = System.getProperties(); + props.setProperty("com.amazonaws.xray.traceHeader", "sys-prop-345"); + + try { + TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor(); + ExecutionAttributes executionAttributes = new ExecutionAttributes(); + + interceptor.beforeExecution(null, executionAttributes); + + Context.ModifyHttpRequest context = context(); + SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes); + + assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + } finally { + SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID"); + props.remove("com.amazonaws.xray.traceHeader"); + } + }); + } + + @Test + public void modifyHttpRequest_whenNotInLambdaEnvironmentWithInternalThreadLocal_shouldNotAddHeader() { + EnvironmentVariableHelper.run(env -> { + resetRelevantEnvVars(env); + + SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "should-be-ignored"); + + try { + TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor(); + ExecutionAttributes executionAttributes = new ExecutionAttributes(); + + interceptor.beforeExecution(null, executionAttributes); + + Context.ModifyHttpRequest context = context(); + SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes); + + assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).isEmpty(); + } finally { + SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID"); + } + }); + } + private Context.ModifyHttpRequest context() { return context(SdkHttpRequest.builder() .uri(URI.create("https://localhost")) diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java index 3299e26ef876..a0747444292d 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java @@ -17,17 +17,25 @@ import static org.assertj.core.api.Assertions.assertThat; +import java.util.List; import org.junit.jupiter.api.Test; import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; import software.amazon.awssdk.awscore.interceptor.TraceIdExecutionInterceptor; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; import software.amazon.awssdk.http.AbortableInputStream; import software.amazon.awssdk.http.HttpExecuteResponse; +import software.amazon.awssdk.http.SdkHttpRequest; import software.amazon.awssdk.http.SdkHttpResponse; import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient; import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient; import software.amazon.awssdk.testutils.EnvironmentVariableHelper; +import software.amazon.awssdk.testutils.service.http.MockAsyncHttpClient; import software.amazon.awssdk.testutils.service.http.MockSyncHttpClient; import software.amazon.awssdk.utils.StringInputStream; +import software.amazon.awssdk.utilslite.SdkInternalThreadLocal; /** * Verifies that the {@link TraceIdExecutionInterceptor} is actually wired up for AWS services. @@ -56,4 +64,181 @@ public void traceIdInterceptorIsEnabled() { } }); } -} + + @Test + public void traceIdInterceptorPreservesTraceIdAcrossRetries() { + EnvironmentVariableHelper.run(env -> { + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123"); + + try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient(); + ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder() + .region(Region.US_WEST_2) + .credentialsProvider(AnonymousCredentialsProvider.create()) + .httpClient(mockHttpClient) + .build()) { + + mockHttpClient.stubResponses( + HttpExecuteResponse.builder() + .response(SdkHttpResponse.builder().statusCode(500).build()) + .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) + .build(), + HttpExecuteResponse.builder() + .response(SdkHttpResponse.builder().statusCode(500).build()) + .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) + .build(), + HttpExecuteResponse.builder().response(SdkHttpResponse.builder().statusCode(200).build()) + .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) + .build()); + + client.allTypes().join(); + + List requests = mockHttpClient.getRequests(); + assertThat(requests).hasSize(3); + + assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + assertThat(requests.get(2).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + + } finally { + SdkInternalThreadLocal.clear(); + } + }); + } + + @Test + public void traceIdInterceptorPreservesTraceIdAcrossChainedFutures() { + EnvironmentVariableHelper.run(env -> { + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123"); + + try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient(); + ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder() + .region(Region.US_WEST_2) + .credentialsProvider(AnonymousCredentialsProvider.create()) + .httpClient(mockHttpClient) + .build()) { + + mockHttpClient.stubResponses( + HttpExecuteResponse.builder() + .response(SdkHttpResponse.builder().statusCode(200).build()) + .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) + .build(), + HttpExecuteResponse.builder() + .response(SdkHttpResponse.builder().statusCode(200).build()) + .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) + .build() + ); + + client.allTypes() + .thenRun(() -> { + client.allTypes().join(); + }) + .join(); + + List requests = mockHttpClient.getRequests(); + + assertThat(requests).hasSize(2); + + assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + + } finally { + SdkInternalThreadLocal.clear(); + } + }); + } + + @Test + public void traceIdInterceptorPreservesTraceIdAcrossExceptionallyCompletedFutures() { + EnvironmentVariableHelper.run(env -> { + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123"); + + try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient(); + ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder() + .region(Region.US_WEST_2) + .credentialsProvider(AnonymousCredentialsProvider.create()) + .httpClient(mockHttpClient) + .build()) { + + mockHttpClient.stubResponses( + HttpExecuteResponse.builder() + .response(SdkHttpResponse.builder().statusCode(400).build()) + .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) + .build(), + HttpExecuteResponse.builder() + .response(SdkHttpResponse.builder().statusCode(200).build()) + .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) + .build() + ); + + client.allTypes() + .exceptionally(throwable -> { + client.allTypes().join(); + return null; + }).join(); + + List requests = mockHttpClient.getRequests(); + + assertThat(requests).hasSize(2); + + assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + + } finally { + SdkInternalThreadLocal.clear(); + } + }); + } + + @Test + public void traceIdInterceptorPreservesTraceIdAcrossExceptionallyCompletedFuturesThrownInPreExecution() { + EnvironmentVariableHelper.run(env -> { + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123"); + + ExecutionInterceptor throwingInterceptor = new ExecutionInterceptor() { + private boolean hasThrown = false; + + @Override + public void beforeMarshalling(Context.BeforeMarshalling context, ExecutionAttributes executionAttributes) { + if (!hasThrown) { + hasThrown = true; + throw new RuntimeException("failing in pre execution"); + } + } + }; + + try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient(); + ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder() + .region(Region.US_WEST_2) + .credentialsProvider(AnonymousCredentialsProvider.create()) + .overrideConfiguration(o -> o.addExecutionInterceptor(throwingInterceptor)) + .httpClient(mockHttpClient) + .build()) { + + mockHttpClient.stubResponses( + HttpExecuteResponse.builder() + .response(SdkHttpResponse.builder().statusCode(200).build()) + .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) + .build() + ); + + client.allTypes() + .exceptionally(throwable -> { + client.allTypes().join(); + return null; + }).join(); + + List requests = mockHttpClient.getRequests(); + + assertThat(requests).hasSize(1); + assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + + } finally { + SdkInternalThreadLocal.clear(); + } + }); + } +} \ No newline at end of file From ea087e25c5917bf2b66077f91e8b4b98083127df Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Fri, 5 Sep 2025 11:02:46 -0700 Subject: [PATCH 03/20] Add changelog --- .changes/next-release/feature-AWSSDKforJavav2-a3984eb.json | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 .changes/next-release/feature-AWSSDKforJavav2-a3984eb.json diff --git a/.changes/next-release/feature-AWSSDKforJavav2-a3984eb.json b/.changes/next-release/feature-AWSSDKforJavav2-a3984eb.json new file mode 100644 index 000000000000..b4336d4002d7 --- /dev/null +++ b/.changes/next-release/feature-AWSSDKforJavav2-a3984eb.json @@ -0,0 +1,6 @@ +{ + "type": "feature", + "category": "AWS SDK for Java v2", + "contributor": "", + "description": "Add a utils-lite package that wraps threadlocal meant for internal shared usage across multiple applications across multiple applications" +} From 77fd847d173ae0fb481587b6bbf0580eb06c234b Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Fri, 5 Sep 2025 11:30:38 -0700 Subject: [PATCH 04/20] update utils-lite version to match parent --- utils-lite/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils-lite/pom.xml b/utils-lite/pom.xml index 9642803ccd58..49da28ec403b 100644 --- a/utils-lite/pom.xml +++ b/utils-lite/pom.xml @@ -21,7 +21,7 @@ software.amazon.awssdk aws-sdk-java-pom - 2.33.3-SNAPSHOT + 2.33.4-SNAPSHOT utils-lite AWS Java SDK :: Utils Lite From 2c6749a773b46b36280b06e7580567af450b40c8 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Fri, 5 Sep 2025 14:32:23 -0700 Subject: [PATCH 05/20] Fix again version --- utils-lite/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils-lite/pom.xml b/utils-lite/pom.xml index 49da28ec403b..9642803ccd58 100644 --- a/utils-lite/pom.xml +++ b/utils-lite/pom.xml @@ -21,7 +21,7 @@ software.amazon.awssdk aws-sdk-java-pom - 2.33.4-SNAPSHOT + 2.33.3-SNAPSHOT utils-lite AWS Java SDK :: Utils Lite From 91422dc71a00e850993128eca13cdd13bd07252f Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Thu, 11 Sep 2025 10:58:20 -0700 Subject: [PATCH 06/20] Bump version --- utils-lite/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils-lite/pom.xml b/utils-lite/pom.xml index 9642803ccd58..2cc7381ee0f0 100644 --- a/utils-lite/pom.xml +++ b/utils-lite/pom.xml @@ -21,7 +21,7 @@ software.amazon.awssdk aws-sdk-java-pom - 2.33.3-SNAPSHOT + 2.33.8-SNAPSHOT utils-lite AWS Java SDK :: Utils Lite From a3be25338b9479a3592b4da05ef93feaa0b43757 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Thu, 11 Sep 2025 13:25:28 -0700 Subject: [PATCH 07/20] Fix braziljson name --- .brazil.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.brazil.json b/.brazil.json index 362888889178..ed43716ef601 100644 --- a/.brazil.json +++ b/.brazil.json @@ -32,9 +32,9 @@ "s3-transfer-manager": { "packageName": "AwsJavaSdk-S3-TransferManager" }, "s3-event-notifications": { "packageName": "AwsJavaSdk-S3-EventNotifications" }, "sdk-core": { "packageName": "AwsJavaSdk-Core" }, - "utils-lite": { "packageName": "AwsJavaSdk-UtilsLite" }, "url-connection-client": { "packageName": "AwsJavaSdk-HttpClient-UrlConnectionClient" }, "utils": { "packageName": "AwsJavaSdk-Core-Utils" }, + "utils-lite": { "packageName": "AwsJavaSdk-Core-UtilsLite" }, "imds": { "packageName": "AwsJavaSdk-Imds" }, "crt-core": { "packageName": "AwsJavaSdk-Core-CrtCore" }, "checksums-spi": { "packageName": "AwsJavaSdk-Core-ChecksumsSpi" }, From e19bd55a21de85acc387c9965f9480322b24f4da Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Thu, 11 Sep 2025 15:45:10 -0700 Subject: [PATCH 08/20] Add thread safe annotation --- .../amazon/awssdk/utilslite/SdkInternalThreadLocal.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java index 34f3b5c9df30..8e7d5b426733 100644 --- a/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java +++ b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java @@ -18,10 +18,12 @@ import java.util.HashMap; import java.util.Map; import software.amazon.awssdk.annotations.SdkProtectedApi; +import software.amazon.awssdk.annotations.ThreadSafe; /** * Utility for thread-local context storage. */ +@ThreadSafe @SdkProtectedApi public final class SdkInternalThreadLocal { private static final ThreadLocal> STORAGE = ThreadLocal.withInitial(HashMap::new); From 71dd50079bf41223cc06dd5c3e7a20069b6c3b7d Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Thu, 11 Sep 2025 16:58:46 -0700 Subject: [PATCH 09/20] Adding a lazy loading of the hashmap --- .../utilslite/SdkInternalThreadLocal.java | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java index 8e7d5b426733..d9f979c978b1 100644 --- a/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java +++ b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java @@ -26,32 +26,44 @@ @ThreadSafe @SdkProtectedApi public final class SdkInternalThreadLocal { - private static final ThreadLocal> STORAGE = ThreadLocal.withInitial(HashMap::new); + private static final ThreadLocal> STORAGE = new ThreadLocal<>(); private SdkInternalThreadLocal() { } public static void put(String key, String value) { + Map map = STORAGE.get(); + if (map == null) { + map = new HashMap<>(); + STORAGE.set(map); + } + if (value == null) { - STORAGE.get().remove(key); + map.remove(key); } else { - STORAGE.get().put(key, value); + map.put(key, value); } } public static String get(String key) { - return STORAGE.get().get(key); + Map map = STORAGE.get(); + return map != null ? map.get(key) : null; } public static String remove(String key) { - return STORAGE.get().remove(key); + Map map = STORAGE.get(); + return map != null ? map.remove(key) : null; } public static void clear() { - STORAGE.get().clear(); + Map map = STORAGE.get(); + if (map != null) { + map.clear(); + } } public static boolean containsKey(String key) { - return STORAGE.get().containsKey(key); + Map map = STORAGE.get(); + return map != null && map.containsKey(key); } } \ No newline at end of file From 6ab8337cf80edb53a6ba6a6ed25247640fecd4a6 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Fri, 12 Sep 2025 10:12:02 -0700 Subject: [PATCH 10/20] Refactored code according to comments --- .../TraceIdExecutionInterceptor.java | 2 +- .../utilslite/SdkInternalThreadLocal.java | 33 +++++++++++++------ 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java index d2b422c940eb..eac667d90be7 100644 --- a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java +++ b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java @@ -28,7 +28,7 @@ /** * The {@code TraceIdExecutionInterceptor} copies the trace details to the {@link #TRACE_ID_HEADER} header, assuming we seem to - * be running in a lambda environment. + * be running in a lambda environment.` */ @SdkProtectedApi public class TraceIdExecutionInterceptor implements ExecutionInterceptor { diff --git a/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java index d9f979c978b1..0e90b6f14ede 100644 --- a/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java +++ b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java @@ -32,37 +32,50 @@ private SdkInternalThreadLocal() { } public static void put(String key, String value) { - Map map = STORAGE.get(); - if (map == null) { - map = new HashMap<>(); - STORAGE.set(map); + if (key == null) { + throw new IllegalArgumentException("Key cannot be null"); } - + Map map = STORAGE.get(); if (value == null) { - map.remove(key); + if (map != null) { + map.remove(key); + if (map.isEmpty()) { + STORAGE.remove(); + } + } } else { + if (map == null) { + map = new HashMap<>(); + STORAGE.set(map); + } map.put(key, value); } } public static String get(String key) { + if (key == null) { + throw new IllegalArgumentException("Key cannot be null"); + } Map map = STORAGE.get(); return map != null ? map.get(key) : null; } public static String remove(String key) { + if (key == null) { + throw new IllegalArgumentException("Key cannot be null"); + } Map map = STORAGE.get(); return map != null ? map.remove(key) : null; } public static void clear() { - Map map = STORAGE.get(); - if (map != null) { - map.clear(); - } + STORAGE.remove(); } public static boolean containsKey(String key) { + if (key == null) { + throw new IllegalArgumentException("Key cannot be null"); + } Map map = STORAGE.get(); return map != null && map.containsKey(key); } From bd8bcf6c1ca03a3fb9c722ac2f498942b25cccbb Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Fri, 12 Sep 2025 12:56:38 -0700 Subject: [PATCH 11/20] Modify threadLocal to InheritableThreadLocal --- .../amazon/awssdk/services/TraceIdTest.java | 34 +++++++++++++++++++ utils-lite/pom.xml | 2 +- .../utilslite/SdkInternalThreadLocal.java | 2 +- 3 files changed, 36 insertions(+), 2 deletions(-) diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java index a0747444292d..8b66e0c32361 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java @@ -241,4 +241,38 @@ public void beforeMarshalling(Context.BeforeMarshalling context, ExecutionAttrib } }); } + + @Test + public void traceIdInterceptorWithNewThreadInheritsTraceId() 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()); + + Thread childThread = new Thread(client::allTypes); + childThread.start(); + childThread.join(); + + List requests = mockHttpClient.getRequests(); + assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123"); + + } catch (InterruptedException e) { + throw new RuntimeException(e); + } finally { + SdkInternalThreadLocal.clear(); + } + }); + } } \ No newline at end of file diff --git a/utils-lite/pom.xml b/utils-lite/pom.xml index 2cc7381ee0f0..08bdeeb52592 100644 --- a/utils-lite/pom.xml +++ b/utils-lite/pom.xml @@ -21,7 +21,7 @@ software.amazon.awssdk aws-sdk-java-pom - 2.33.8-SNAPSHOT + 2.33.9-SNAPSHOT utils-lite AWS Java SDK :: Utils Lite diff --git a/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java index 0e90b6f14ede..d4d0dc4011a4 100644 --- a/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java +++ b/utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java @@ -26,7 +26,7 @@ @ThreadSafe @SdkProtectedApi public final class SdkInternalThreadLocal { - private static final ThreadLocal> STORAGE = new ThreadLocal<>(); + private static final ThreadLocal> STORAGE = new InheritableThreadLocal<>(); private SdkInternalThreadLocal() { } From e576b98fd1a09b586c3741f555219cdb4f85e7af Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Sun, 14 Sep 2025 20:01:38 -0700 Subject: [PATCH 12/20] Add javadoc clarification --- .../awscore/interceptor/TraceIdExecutionInterceptor.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java index eac667d90be7..b4192e00d103 100644 --- a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java +++ b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java @@ -69,6 +69,10 @@ public void onExecutionFailure(Context.FailedExecution context, ExecutionAttribu saveTraceId(executionAttributes); } + /** + * Stores the trace ID in thread-local storage to ensure trace propagation across + * thread boundaries during retries, or future chaining. + */ private static void saveTraceId(ExecutionAttributes executionAttributes) { String traceId = executionAttributes.getAttribute(TRACE_ID); if (traceId != null) { From fa81619860b250913c068e5dc72cde2ab47d3d3d Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Sun, 14 Sep 2025 20:06:38 -0700 Subject: [PATCH 13/20] Add lambda check to prevent extra checks on non-lambda environments --- .../interceptor/TraceIdExecutionInterceptor.java | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java index b4192e00d103..0e9ad0ff9a13 100644 --- a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java +++ b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java @@ -40,7 +40,7 @@ public class TraceIdExecutionInterceptor implements ExecutionInterceptor { @Override public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) { String traceId = SdkInternalThreadLocal.get(CONCURRENT_TRACE_ID_KEY); - if (traceId != null) { + if (traceId != null && lambdaFunctionNameEnvironmentVariable().isPresent()) { executionAttributes.putAttribute(TRACE_ID, traceId); } } @@ -61,12 +61,16 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, Execu @Override public void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) { - saveTraceId(executionAttributes); + if (lambdaFunctionNameEnvironmentVariable().isPresent()) { + saveTraceId(executionAttributes); + } } @Override public void onExecutionFailure(Context.FailedExecution context, ExecutionAttributes executionAttributes) { - saveTraceId(executionAttributes); + if (lambdaFunctionNameEnvironmentVariable().isPresent()) { + saveTraceId(executionAttributes); + } } /** From c4294823c00dfa597986d8f247a7a629afb526a9 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Mon, 15 Sep 2025 12:54:06 -0700 Subject: [PATCH 14/20] Update beforeExecution logic --- .../awscore/interceptor/TraceIdExecutionInterceptor.java | 8 +++++--- utils-lite/pom.xml | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java index 0e9ad0ff9a13..5f686e793b5f 100644 --- a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java +++ b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java @@ -39,9 +39,11 @@ public class TraceIdExecutionInterceptor implements ExecutionInterceptor { @Override public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) { - String traceId = SdkInternalThreadLocal.get(CONCURRENT_TRACE_ID_KEY); - if (traceId != null && lambdaFunctionNameEnvironmentVariable().isPresent()) { - executionAttributes.putAttribute(TRACE_ID, traceId); + if(lambdaFunctionNameEnvironmentVariable().isPresent()) { + String traceId = SdkInternalThreadLocal.get(CONCURRENT_TRACE_ID_KEY); + if (traceId != null) { + executionAttributes.putAttribute(TRACE_ID, traceId); + } } } diff --git a/utils-lite/pom.xml b/utils-lite/pom.xml index 08bdeeb52592..06632e9812f3 100644 --- a/utils-lite/pom.xml +++ b/utils-lite/pom.xml @@ -21,7 +21,7 @@ software.amazon.awssdk aws-sdk-java-pom - 2.33.9-SNAPSHOT + 2.33.10-SNAPSHOT utils-lite AWS Java SDK :: Utils Lite From bf72bd6a8768285fc0fb441a84ae7671af3781ac Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Mon, 15 Sep 2025 14:29:10 -0700 Subject: [PATCH 15/20] Fix checkstyle --- .../awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java index 5f686e793b5f..0b24d55bd0e6 100644 --- a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java +++ b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java @@ -39,7 +39,7 @@ public class TraceIdExecutionInterceptor implements ExecutionInterceptor { @Override public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) { - if(lambdaFunctionNameEnvironmentVariable().isPresent()) { + if (lambdaFunctionNameEnvironmentVariable().isPresent()) { String traceId = SdkInternalThreadLocal.get(CONCURRENT_TRACE_ID_KEY); if (traceId != null) { executionAttributes.putAttribute(TRACE_ID, traceId); From 23a1e04dfa11faea7835ddcdc12ac7b2f49535d5 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Wed, 17 Sep 2025 09:43:56 -0700 Subject: [PATCH 16/20] Add test coverage --- aws-sdk-java/pom.xml | 5 ++ .../amazon/awssdk/services/TraceIdTest.java | 69 ++++++++++++++++++- 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/aws-sdk-java/pom.xml b/aws-sdk-java/pom.xml index ee28fa8df3c8..0e3b605c89b7 100644 --- a/aws-sdk-java/pom.xml +++ b/aws-sdk-java/pom.xml @@ -2113,6 +2113,11 @@ Amazon AutoScaling, etc). bcmdashboards ${awsjavasdk.version} + + software.amazon.awssdk + utils-lite + ${awsjavasdk.version} + ${project.artifactId}-${project.version} diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java index 8b66e0c32361..5d4a167a0efb 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java +++ b/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 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 requests = mockHttpClient.getRequests(); + assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).isEmpty(); + + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } finally { + SdkInternalThreadLocal.clear(); + } + }); + } } \ No newline at end of file From 445184ec8228c758d5f15217e035b04a814c95ac Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Wed, 17 Sep 2025 10:03:43 -0700 Subject: [PATCH 17/20] update version --- utils-lite/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils-lite/pom.xml b/utils-lite/pom.xml index 06632e9812f3..51b3f680d301 100644 --- a/utils-lite/pom.xml +++ b/utils-lite/pom.xml @@ -21,7 +21,7 @@ software.amazon.awssdk aws-sdk-java-pom - 2.33.10-SNAPSHOT + 2.33.12-SNAPSHOT utils-lite AWS Java SDK :: Utils Lite From c96858ed41647ef7d9e1b2126544e308ede34932 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Wed, 17 Sep 2025 13:38:45 -0700 Subject: [PATCH 18/20] Removing non-determinstic test --- .../amazon/awssdk/services/TraceIdTest.java | 32 +------------------ 1 file changed, 1 insertion(+), 31 deletions(-) diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java index 5d4a167a0efb..415fdc5d1813 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java @@ -22,6 +22,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import org.junit.jupiter.api.RepeatedTest; import org.junit.jupiter.api.Test; import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; import software.amazon.awssdk.awscore.interceptor.TraceIdExecutionInterceptor; @@ -311,35 +312,4 @@ public void traceIdInterceptorWithExecutiveServicePreservesTraceId() { } }); } - - @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 requests = mockHttpClient.getRequests(); - assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).isEmpty(); - - } catch (InterruptedException | ExecutionException e) { - throw new RuntimeException(e); - } finally { - SdkInternalThreadLocal.clear(); - } - }); - } } \ No newline at end of file From 5bcca7094827da0c2a122f0180c5c719850fda0f Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Wed, 17 Sep 2025 13:46:20 -0700 Subject: [PATCH 19/20] Removing unused dependencies --- .../test/java/software/amazon/awssdk/services/TraceIdTest.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java index 415fdc5d1813..87de50ec03ae 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java @@ -18,11 +18,9 @@ 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.RepeatedTest; import org.junit.jupiter.api.Test; import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; import software.amazon.awssdk.awscore.interceptor.TraceIdExecutionInterceptor; From 73427bb0a9a21161d20cbe7615fabfd44c690f99 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Thu, 18 Sep 2025 09:56:30 -0700 Subject: [PATCH 20/20] update version --- utils-lite/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils-lite/pom.xml b/utils-lite/pom.xml index 51b3f680d301..e725607bb44a 100644 --- a/utils-lite/pom.xml +++ b/utils-lite/pom.xml @@ -21,7 +21,7 @@ software.amazon.awssdk aws-sdk-java-pom - 2.33.12-SNAPSHOT + 2.33.13-SNAPSHOT utils-lite AWS Java SDK :: Utils Lite