Add ThreadStorage utility class and support for Lambda TraceId propagation

.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" },

.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."
+}

aws-sdk-java/pom.xml

@@ -813,6 +813,11 @@ Amazon AutoScaling, etc).</description>
             <artifactId>swf</artifactId>
             <version>${awsjavasdk.version}</version>
         </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>utils-lite</artifactId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>textract</artifactId>

bom/pom.xml

@@ -252,6 +252,11 @@
                 <artifactId>endpoints-spi</artifactId>
                 <version>${awsjavasdk.version}</version>
             </dependency>
+            <dependency>
+                <groupId>software.amazon.awssdk</groupId>
+                <artifactId>utils-lite</artifactId>
+                <version>${awsjavasdk.version}</version>
+            </dependency>
             <!-- Services -->
             <dependency>
                 <groupId>software.amazon.awssdk</groupId>

core/aws-core/pom.xml

@@ -113,7 +113,11 @@
             <groupId>software.amazon.eventstream</groupId>
             <artifactId>eventstream</artifactId>
         </dependency>
-
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>utils-lite</artifactId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>test-utils</artifactId>

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<String> 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<String> traceIdHeader = traceIdHeader(context);
         if (!traceIdHeader.isPresent()) {
             Optional<String> lambdafunctionName = lambdaFunctionNameEnvironmentVariable();
-            Optional<String> traceId = traceId();
+            Optional<String> 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<String> traceIdHeader(Context.ModifyHttpRequest context) {
         return context.httpRequest().firstMatchingHeader(TRACE_ID_HEADER);
     }
 
-    private Optional<String> traceId() {
+    private Optional<String> traceId(ExecutionAttributes executionAttributes) {
+        Optional<String> 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<String> lambdaFunctionNameEnvironmentVariable() {
         return SystemSetting.getStringValueFromEnvironmentVariable(LAMBDA_FUNCTION_NAME_ENVIRONMENT_VARIABLE);
         // CHECKSTYLE:ON
     }
-}
+}

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"))

pom.xml

@@ -67,6 +67,7 @@
         <module>metric-publishers</module>
         <module>release-scripts</module>
         <module>utils</module>
+        <module>utils-lite</module>
         <module>codegen-lite</module>
         <module>codegen-lite-maven-plugin</module>
         <module>archetypes</module>
@@ -665,6 +666,7 @@
                             <includeModule>cloudwatch-metric-publisher</includeModule>
                             <includeModule>emf-metric-logging-publisher</includeModule>
                             <includeModule>utils</includeModule>
+                            <includeModule>utils-lite</includeModule>
                             <includeModule>imds</includeModule>
                             <includeModule>retries</includeModule>
                             <includeModule>retries-spi</includeModule>

test/architecture-tests/pom.xml

@@ -61,6 +61,11 @@
             <groupId>software.amazon.awssdk</groupId>
             <version>${awsjavasdk.version}</version>
         </dependency>
+        <dependency>
+            <artifactId>utils-lite</artifactId>
+            <groupId>software.amazon.awssdk</groupId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
         <dependency>
             <artifactId>s3</artifactId>
             <groupId>software.amazon.awssdk</groupId>

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);
+    }
+}