MultiConcurrent XRAY TraceID using utils-lite (#96)

Co-authored-by: Mohammed Ehab <moehabe@amazon.com>

Author: M-Elsaeed

aws-lambda-java-runtime-interface-client/pom.xml

@@ -1,6 +1,6 @@
 <project xmlns="http://maven.apache.org/POM/4.0.0"
-         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
     <modelVersion>4.0.0</modelVersion>
     <groupId>com.amazonaws</groupId>
     <artifactId>aws-lambda-java-runtime-interface-client</artifactId>
@@ -47,9 +47,9 @@
             separately from the Runtime Interface Client functionality until we figure something else out.
         -->
         <multiArch>true</multiArch>
-        <target_build_os/>
-        <target_build_arch/>
-        <ric.classifier/>
+        <target_build_os />
+        <target_build_arch />
+        <ric.classifier />
         <!--
             Some test/integration/codebuild/buildspec.*.yml files will set -DargLineForReflectionTestOnly=""
             as this is not supported by all flavors of java
@@ -68,7 +68,11 @@
             <artifactId>aws-lambda-java-serialization</artifactId>
             <version>1.1.6</version>
         </dependency>
-
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>utils-lite</artifactId>
+            <version>2.34.0</version>
+        </dependency>
         <dependency>
             <groupId>org.junit.jupiter</groupId>
             <artifactId>junit-jupiter-engine</artifactId>
@@ -152,11 +156,11 @@
                         <configuration>
                             <target name="Build JNI libraries for tests">
                                 <exec executable="${project.basedir}/src/main/jni/build-jni-lib.sh"
-                                      failonerror="true" logError="true">
-                                    <arg value="${project.build.directory}"/>
-                                    <arg value="false"/>
-                                    <arg value="${target_build_os}"/>
-                                    <arg value="${target_build_arch}"/>
+                                    failonerror="true" logError="true">
+                                    <arg value="${project.build.directory}" />
+                                    <arg value="false" />
+                                    <arg value="${target_build_os}" />
+                                    <arg value="${target_build_arch}" />
                                 </exec>
                             </target>
                         </configuration>
@@ -170,11 +174,11 @@
                         <configuration>
                             <target name="Build JNI libraries">
                                 <exec executable="${project.basedir}/src/main/jni/build-jni-lib.sh"
-                                      failonerror="true" logError="true">
-                                    <arg value="${project.build.directory}"/>
-                                    <arg value="${multiArch}"/>
-                                    <arg value="${target_build_os}"/>
-                                    <arg value="${target_build_arch}"/>
+                                    failonerror="true" logError="true">
+                                    <arg value="${project.build.directory}" />
+                                    <arg value="${multiArch}" />
+                                    <arg value="${target_build_os}" />
+                                    <arg value="${target_build_arch}" />
                                 </exec>
                             </target>
                         </configuration>
@@ -269,7 +273,7 @@
                     </execution>
                 </executions>
             </plugin>
-              <plugin>
+            <plugin>
                 <groupId>org.apache.maven.plugins</groupId>
                 <artifactId>maven-checkstyle-plugin</artifactId>
                 <version>${maven-checkstyle-plugin.version}</version>

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/AWSLambda.java

@@ -40,6 +40,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import software.amazon.awssdk.utilslite.SdkInternalThreadLocal;
 
 /**
  * The entrypoint of this class is {@link AWSLambda#startRuntime}. It performs two main tasks:
@@ -73,6 +74,8 @@ public class AWSLambda {
     private static final String INIT_TYPE_SNAP_START = "snap-start";
 
     private static final String AWS_LAMBDA_INITIALIZATION_TYPE = System.getenv(ReservedRuntimeEnvironmentVariables.AWS_LAMBDA_INITIALIZATION_TYPE);
+
+    private static final String CONCURRENT_TRACE_ID_KEY = "AWS_LAMBDA_X_TRACE_ID";
 
     static {
         // Override the disabledAlgorithms setting to match configuration for openjdk8-u181.
@@ -304,7 +307,11 @@ private static void startRuntimeLoop(LambdaRequestHandler lambdaRequestHandler,
             try {
                 UserFault userFault = null;
                 InvocationRequest request = exitLoopOnErrors ? runtimeClient.nextInvocation() : runtimeClient.nextInvocationWithExponentialBackoff(lambdaLogger);
-                setEnvVarForXrayTraceId(request);
+                if (exitLoopOnErrors) {
+                    setEnvVarForXrayTraceId(request);
+                } else {
+                    SdkInternalThreadLocal.put(CONCURRENT_TRACE_ID_KEY, request.getXrayTraceId());
+                }
 
                 try {
                     ByteArrayOutputStream payload = lambdaRequestHandler.call(request);
@@ -321,6 +328,8 @@ private static void startRuntimeLoop(LambdaRequestHandler lambdaRequestHandler,
                     if (userFault != null) {
                         lambdaLogger.log(userFault.reportableError(), lambdaLogger.getLogFormat() == LogFormat.JSON ? LogLevel.ERROR : LogLevel.UNDEFINED);
                     }
+
+                    SdkInternalThreadLocal.remove(CONCURRENT_TRACE_ID_KEY);
                 }
             } catch (Throwable t) {
                 if (exitLoopOnErrors || t instanceof LambdaRuntimeClientMaxRetriesExceededException) {