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/.changes/next-release/feature-AWSSDKforJavav2-2e7c0a3.json b/.changes/next-release/feature-AWSSDKforJavav2-2e7c0a3.json new file mode 100644 index 000000000000..4b0c2020e6f1 --- /dev/null +++ b/.changes/next-release/feature-AWSSDKforJavav2-2e7c0a3.json @@ -0,0 +1,6 @@ +{ + "type": "feature", + "category": "AWS SDK for Java v2", + "contributor": "", + "description": "Adding a small utility class to store data on thread local that can be used across components." +} diff --git a/aws-sdk-java/pom.xml b/aws-sdk-java/pom.xml index 81f8e827a9ff..debe1a4f95a0 100644 --- a/aws-sdk-java/pom.xml +++ b/aws-sdk-java/pom.xml @@ -813,6 +813,11 @@ Amazon AutoScaling, etc). swf ${awsjavasdk.version} + + software.amazon.awssdk + utils-lite + ${awsjavasdk.version} + software.amazon.awssdk textract diff --git a/bom/pom.xml b/bom/pom.xml index ede5b9251007..47ddb4196061 100644 --- a/bom/pom.xml +++ b/bom/pom.xml @@ -252,6 +252,11 @@ endpoints-spi ${awsjavasdk.version} + + software.amazon.awssdk + utils-lite + ${awsjavasdk.version} + software.amazon.awssdk diff --git a/core/aws-core/pom.xml b/core/aws-core/pom.xml index fb01b175d288..82514e2c01c7 100644 --- a/core/aws-core/pom.xml +++ b/core/aws-core/pom.xml @@ -113,7 +113,11 @@ software.amazon.eventstream eventstream - + + software.amazon.awssdk + utils-lite + ${awsjavasdk.version} + software.amazon.awssdk test-utils 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/pom.xml b/pom.xml index 42dbdbfc01bb..f2a7f99dd05f 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 @@ -665,6 +666,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 7b255deff4d5..429b732c0f10 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..b0f8dc5ad1e6 --- /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); + } +} 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 diff --git a/test/http-client-benchmarks/pom.xml b/test/http-client-benchmarks/pom.xml index 10fb40969cc4..5b9cbe70b84b 100644 --- a/test/http-client-benchmarks/pom.xml +++ b/test/http-client-benchmarks/pom.xml @@ -175,7 +175,7 @@ org.apache.maven.plugins maven-shade-plugin - 2.2 + 3.5.0 package diff --git a/test/sdk-benchmarks/pom.xml b/test/sdk-benchmarks/pom.xml index b6d6af022454..736e5001110e 100644 --- a/test/sdk-benchmarks/pom.xml +++ b/test/sdk-benchmarks/pom.xml @@ -340,7 +340,7 @@ org.apache.maven.plugins maven-shade-plugin - 2.2 + 3.5.0 package diff --git a/test/tests-coverage-reporting/pom.xml b/test/tests-coverage-reporting/pom.xml index 39e84897ced9..6cb40d401a44 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..02de38eac13b --- /dev/null +++ b/utils-lite/pom.xml @@ -0,0 +1,79 @@ + + + + + 4.0.0 + + software.amazon.awssdk + aws-sdk-java-pom + 2.33.2-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..3bd84c96d57f --- /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); + } +} 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..42044861cf33 --- /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