Add support for lambda traces for concurrent environments

Author: RanVaknin

core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.awscore.interceptor;
 
 import java.util.Optional;
+import org.slf4j.MDC;
 import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.awscore.internal.interceptor.TracingSystemSetting;
 import software.amazon.awssdk.core.interceptor.Context;
@@ -32,6 +33,7 @@
 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_TraceId";
 
     @Override
     public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
@@ -53,6 +55,9 @@ private Optional<String> traceIdHeader(Context.ModifyHttpRequest context) {
     }
 
     private Optional<String> traceId() {
+        if (TracingSystemSetting.AWS_LAMBDA_MAX_CONCURRENCY.getStringValue().isPresent()) {
+            return Optional.ofNullable(MDC.get(CONCURRENT_TRACE_ID_KEY));
+        }
         return TracingSystemSetting._X_AMZN_TRACE_ID.getStringValue();
     }
 

core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/interceptor/TracingSystemSetting.java

@@ -24,7 +24,9 @@
 @SdkInternalApi
 public enum TracingSystemSetting implements SystemSetting {
     // See: https://github.com/aws/aws-xray-sdk-java/issues/251
-    _X_AMZN_TRACE_ID("com.amazonaws.xray.traceHeader", null);
+    _X_AMZN_TRACE_ID("com.amazonaws.xray.traceHeader", null),
+    // Environment variable to detect Lambda multi concurrency mode ("elevator"). This value is set by the Lambda runtime.
+    AWS_LAMBDA_MAX_CONCURRENCY("aws.lambda.maxConcurrency", null);
 
     private final String systemProperty;
     private final String defaultValue;

core/aws-core/src/test/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptorTest.java

@@ -21,6 +21,7 @@
 import java.util.Properties;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
+import org.slf4j.MDC;
 import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.interceptor.Context;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
@@ -111,6 +112,94 @@ public void headerNotAddedIfNoTraceIdEnvVar() {
         });
     }
 
+    @Test
+    public void modifyHttpRequest_whenMultiConcurrencyModeWithMdc_shouldAddTraceIdHeader() {
+        EnvironmentVariableHelper.run(env -> {
+            resetRelevantEnvVars(env);
+            env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
+            env.set("AWS_LAMBDA_MAX_CONCURRENCY", "10");
+
+            MDC.put("AWS_LAMBDA_X_TraceId", "mdc-trace-123");
+
+            try {
+                Context.ModifyHttpRequest context = context();
+                assertThat(modifyHttpRequest(context).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
+            } finally {
+                MDC.remove("AWS_LAMBDA_X_TraceId");
+            }
+        });
+    }
+
+    @Test
+    public void modifyHttpRequest_whenMultiConcurrencyModeWithBothMdcAndSystemProperty_shouldUseMdcValue() {
+        EnvironmentVariableHelper.run(env -> {
+            resetRelevantEnvVars(env);
+            env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
+            env.set("AWS_LAMBDA_MAX_CONCURRENCY", "10");
+
+            MDC.put("AWS_LAMBDA_X_TraceId", "mdc-trace-123");
+            Properties props = System.getProperties();
+            props.setProperty("com.amazonaws.xray.traceHeader", "sys-prop-345");
+
+            try {
+                Context.ModifyHttpRequest context = context();
+                assertThat(modifyHttpRequest(context).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
+            } finally {
+                MDC.remove("AWS_LAMBDA_X_TraceId");
+                props.remove("com.amazonaws.xray.traceHeader");
+            }
+        });
+    }
+
+    @Test
+    public void modifyHttpRequest_whenMultiConcurrencyModeWithEmptyMdc_shouldNotAddHeader() {
+        EnvironmentVariableHelper.run(env -> {
+            resetRelevantEnvVars(env);
+            env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
+            env.set("AWS_LAMBDA_MAX_CONCURRENCY", "10");
+
+            MDC.clear();
+
+            Context.ModifyHttpRequest context = context();
+            assertThat(modifyHttpRequest(context)).isSameAs(context.httpRequest());
+        });
+    }
+
+    @Test
+    public void modifyHttpRequest_whenNotInLambdaEnvironmentWithMdc_shouldNotAddHeader() {
+        EnvironmentVariableHelper.run(env -> {
+            resetRelevantEnvVars(env);
+            env.set("AWS_LAMBDA_MAX_CONCURRENCY", "10");
+
+            MDC.put("AWS_LAMBDA_X_TraceId", "should-be-ignored");
+
+            try {
+                Context.ModifyHttpRequest context = context();
+                assertThat(modifyHttpRequest(context)).isSameAs(context.httpRequest());
+            } finally {
+                MDC.remove("AWS_LAMBDA_X_TraceId");
+            }
+        });
+    }
+
+    @Test
+    public void modifyHttpRequest_whenConcurrencyModeIsEmptyString_shouldUseMdcValue() {
+        EnvironmentVariableHelper.run(env -> {
+            resetRelevantEnvVars(env);
+            env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
+            env.set("AWS_LAMBDA_MAX_CONCURRENCY", "");
+
+            MDC.put("AWS_LAMBDA_X_TraceId", "empty-string-test");
+
+            try {
+                Context.ModifyHttpRequest context = context();
+                assertThat(modifyHttpRequest(context).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("empty-string-test");
+            } finally {
+                MDC.remove("AWS_LAMBDA_X_TraceId");
+            }
+        });
+    }
+
     private Context.ModifyHttpRequest context() {
         return context(SdkHttpRequest.builder()
                                      .uri(URI.create("https://localhost"))
@@ -133,5 +222,6 @@ private SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context) {
     private void resetRelevantEnvVars(EnvironmentVariableHelper env) {
         env.remove("AWS_LAMBDA_FUNCTION_NAME");
         env.remove("_X_AMZN_TRACE_ID");
+        env.remove("AWS_LAMBDA_MAX_CONCURRENCY");
     }
 }