From 5cfbfb14938bd6475261624e4ce89e0236b49cba Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Wed, 23 Jul 2025 16:03:25 -0700 Subject: [PATCH 1/6] Add support for lambda traces for concurrent environments --- .../TraceIdExecutionInterceptor.java | 5 ++ .../interceptor/TracingSystemSetting.java | 4 +- .../TraceIdExecutionInterceptorTest.java | 90 +++++++++++++++++++ 3 files changed, 98 insertions(+), 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 95224228cfb4..12eb70f501f7 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 @@ -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 traceIdHeader(Context.ModifyHttpRequest context) { } private Optional 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(); } diff --git a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/interceptor/TracingSystemSetting.java b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/interceptor/TracingSystemSetting.java index 6f412e9a83a5..ff6523b7502a 100644 --- a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/interceptor/TracingSystemSetting.java +++ b/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; 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..b1f145c3ad77 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 @@ -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"); } } \ No newline at end of file From e4f8dd25e0ae40cb3a1ae0847bb8d02ffac10e4d Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Fri, 25 Jul 2025 14:10:14 -0700 Subject: [PATCH 2/6] Add async support --- .../TraceIdExecutionInterceptor.java | 24 +++++--- .../interceptor/TracingSystemSetting.java | 4 +- .../TraceIdExecutionInterceptorTest.java | 58 +++++++------------ .../amazon/awssdk/services/TraceIdTest.java | 41 +++++++++++++ 4 files changed, 80 insertions(+), 47 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 12eb70f501f7..cf3c62ed73f1 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 @@ -20,6 +20,7 @@ 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; @@ -34,19 +35,27 @@ 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"; + protected static final ExecutionAttribute CACHED_TRACE_ID = new ExecutionAttribute<>("CachedTraceId"); + + @Override + public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) { + String traceId = MDC.get(CONCURRENT_TRACE_ID_KEY); + if (traceId != null) { + executionAttributes.putAttribute(CACHED_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 lambdaFunctionName = lambdaFunctionNameEnvironmentVariable(); + Optional traceId = traceId(executionAttributes); - if (lambdafunctionName.isPresent() && traceId.isPresent()) { + if (lambdaFunctionName.isPresent() && traceId.isPresent()) { return context.httpRequest().copy(r -> r.putHeader(TRACE_ID_HEADER, traceId.get())); } } - return context.httpRequest(); } @@ -54,9 +63,10 @@ private Optional traceIdHeader(Context.ModifyHttpRequest context) { return context.httpRequest().firstMatchingHeader(TRACE_ID_HEADER); } - private Optional traceId() { - if (TracingSystemSetting.AWS_LAMBDA_MAX_CONCURRENCY.getStringValue().isPresent()) { - return Optional.ofNullable(MDC.get(CONCURRENT_TRACE_ID_KEY)); + private Optional traceId(ExecutionAttributes executionAttributes) { + Optional traceId = Optional.ofNullable(executionAttributes.getAttribute(CACHED_TRACE_ID)); + if (traceId.isPresent()) { + return traceId; } return TracingSystemSetting._X_AMZN_TRACE_ID.getStringValue(); } diff --git a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/interceptor/TracingSystemSetting.java b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/interceptor/TracingSystemSetting.java index ff6523b7502a..6f412e9a83a5 100644 --- a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/interceptor/TracingSystemSetting.java +++ b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/interceptor/TracingSystemSetting.java @@ -24,9 +24,7 @@ @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), - // 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); + _X_AMZN_TRACE_ID("com.amazonaws.xray.traceHeader", null); private final String systemProperty; private final String defaultValue; 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 b1f145c3ad77..875d530f489d 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 @@ -117,13 +117,17 @@ public void modifyHttpRequest_whenMultiConcurrencyModeWithMdc_shouldAddTraceIdHe 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 { + TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor(); + ExecutionAttributes executionAttributes = new ExecutionAttributes(); + + interceptor.beforeExecution(null, executionAttributes); Context.ModifyHttpRequest context = context(); - assertThat(modifyHttpRequest(context).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); + + SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes); + assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); } finally { MDC.remove("AWS_LAMBDA_X_TraceId"); } @@ -135,15 +139,21 @@ public void modifyHttpRequest_whenMultiConcurrencyModeWithBothMdcAndSystemProper 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 { + TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor(); + ExecutionAttributes executionAttributes = new ExecutionAttributes(); + + interceptor.beforeExecution(null, executionAttributes); + Context.ModifyHttpRequest context = context(); - assertThat(modifyHttpRequest(context).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); + SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes); + + assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); } finally { MDC.remove("AWS_LAMBDA_X_TraceId"); props.remove("com.amazonaws.xray.traceHeader"); @@ -151,49 +161,23 @@ public void modifyHttpRequest_whenMultiConcurrencyModeWithBothMdcAndSystemProper }); } - @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"); - } - }); - } + TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor(); + ExecutionAttributes executionAttributes = new ExecutionAttributes(); - @Test - public void modifyHttpRequest_whenConcurrencyModeIsEmptyString_shouldUseMdcValue() { - EnvironmentVariableHelper.run(env -> { - resetRelevantEnvVars(env); - env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); - env.set("AWS_LAMBDA_MAX_CONCURRENCY", ""); + interceptor.beforeExecution(null, executionAttributes); - 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"); + SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes); + + assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).isEmpty(); } finally { MDC.remove("AWS_LAMBDA_X_TraceId"); } 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..4102185ead10 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,15 +17,20 @@ import static org.assertj.core.api.Assertions.assertThat; +import java.util.List; import org.junit.jupiter.api.Test; +import org.slf4j.MDC; import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; import software.amazon.awssdk.awscore.interceptor.TraceIdExecutionInterceptor; 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; @@ -56,4 +61,40 @@ public void traceIdInterceptorIsEnabled() { } }); } + + @Test + public void traceIdInterceptorPreservesTraceIdAcrossRetries() { + EnvironmentVariableHelper.run(env -> { + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + MDC.put("AWS_LAMBDA_X_TraceId", "mdc-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(200).build()) + .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) + .build()); + + client.allTypes().join(); + + List requests = mockHttpClient.getRequests(); + assertThat(requests).hasSize(2); + + assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); + assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); + + } finally { + MDC.clear(); + } + }); + } } From d6967563a8939008332da485c722dc1844502102 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Fri, 25 Jul 2025 15:07:00 -0700 Subject: [PATCH 3/6] fix access modifier --- .../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 cf3c62ed73f1..d4a45e93af18 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 @@ -35,7 +35,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"; - protected static final ExecutionAttribute CACHED_TRACE_ID = new ExecutionAttribute<>("CachedTraceId"); + private static final ExecutionAttribute CACHED_TRACE_ID = new ExecutionAttribute<>("CachedTraceId"); @Override public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) { From e0814cb0601d90d5c42cb7ceafc804ed394d8d45 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Wed, 30 Jul 2025 20:24:37 -0700 Subject: [PATCH 4/6] Add afterExecution/onExecutionFailure hooks to restore MDC trace ID --- .../TraceIdExecutionInterceptor.java | 22 ++++- .../TraceIdExecutionInterceptorTest.java | 13 ++- .../amazon/awssdk/services/TraceIdTest.java | 98 ++++++++++++++++++- 3 files changed, 120 insertions(+), 13 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 d4a45e93af18..a16132a9d834 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 @@ -34,14 +34,14 @@ 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"; - private static final ExecutionAttribute CACHED_TRACE_ID = new ExecutionAttribute<>("CachedTraceId"); + 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 = MDC.get(CONCURRENT_TRACE_ID_KEY); if (traceId != null) { - executionAttributes.putAttribute(CACHED_TRACE_ID, traceId); + executionAttributes.putAttribute(TRACE_ID, traceId); } } @@ -59,12 +59,26 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, Execu 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) { + MDC.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(ExecutionAttributes executionAttributes) { - Optional traceId = Optional.ofNullable(executionAttributes.getAttribute(CACHED_TRACE_ID)); + Optional traceId = Optional.ofNullable(executionAttributes.getAttribute(TRACE_ID)); if (traceId.isPresent()) { return traceId; } 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 875d530f489d..91992716dc2a 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 @@ -117,7 +117,7 @@ public void modifyHttpRequest_whenMultiConcurrencyModeWithMdc_shouldAddTraceIdHe EnvironmentVariableHelper.run(env -> { resetRelevantEnvVars(env); env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); - MDC.put("AWS_LAMBDA_X_TraceId", "mdc-trace-123"); + MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123"); try { TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor(); @@ -129,7 +129,7 @@ public void modifyHttpRequest_whenMultiConcurrencyModeWithMdc_shouldAddTraceIdHe SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes); assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); } finally { - MDC.remove("AWS_LAMBDA_X_TraceId"); + MDC.remove("AWS_LAMBDA_X_TRACE_ID"); } }); } @@ -140,7 +140,7 @@ public void modifyHttpRequest_whenMultiConcurrencyModeWithBothMdcAndSystemProper resetRelevantEnvVars(env); env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); - MDC.put("AWS_LAMBDA_X_TraceId", "mdc-trace-123"); + MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123"); Properties props = System.getProperties(); props.setProperty("com.amazonaws.xray.traceHeader", "sys-prop-345"); @@ -155,7 +155,7 @@ public void modifyHttpRequest_whenMultiConcurrencyModeWithBothMdcAndSystemProper assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); } finally { - MDC.remove("AWS_LAMBDA_X_TraceId"); + MDC.remove("AWS_LAMBDA_X_TRACE_ID"); props.remove("com.amazonaws.xray.traceHeader"); } }); @@ -166,7 +166,7 @@ public void modifyHttpRequest_whenNotInLambdaEnvironmentWithMdc_shouldNotAddHead EnvironmentVariableHelper.run(env -> { resetRelevantEnvVars(env); - MDC.put("AWS_LAMBDA_X_TraceId", "should-be-ignored"); + MDC.put("AWS_LAMBDA_X_TRACE_ID", "should-be-ignored"); try { TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor(); @@ -179,7 +179,7 @@ public void modifyHttpRequest_whenNotInLambdaEnvironmentWithMdc_shouldNotAddHead assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).isEmpty(); } finally { - MDC.remove("AWS_LAMBDA_X_TraceId"); + MDC.remove("AWS_LAMBDA_X_TRACE_ID"); } }); } @@ -206,6 +206,5 @@ 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"); } } \ No newline at end of file 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 4102185ead10..13f799d63e0b 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 @@ -37,7 +37,8 @@ /** * Verifies that the {@link TraceIdExecutionInterceptor} is actually wired up for AWS services. */ -public class TraceIdTest { +public class +TraceIdTest { @Test public void traceIdInterceptorIsEnabled() { EnvironmentVariableHelper.run(env -> { @@ -66,7 +67,7 @@ public void traceIdInterceptorIsEnabled() { public void traceIdInterceptorPreservesTraceIdAcrossRetries() { EnvironmentVariableHelper.run(env -> { env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); - MDC.put("AWS_LAMBDA_X_TraceId", "mdc-trace-123"); + MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123"); try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient(); ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder() @@ -76,6 +77,10 @@ public void traceIdInterceptorPreservesTraceIdAcrossRetries() { .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("{}"))) @@ -87,6 +92,94 @@ public void traceIdInterceptorPreservesTraceIdAcrossRetries() { client.allTypes().join(); List requests = mockHttpClient.getRequests(); + assertThat(requests).hasSize(3); + + assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); + assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); + assertThat(requests.get(2).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); + + } finally { + MDC.clear(); + } + }); + } + + @Test + public void traceIdInterceptorPreservesTraceIdAcrossChainedFutures() { + EnvironmentVariableHelper.run(env -> { + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-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(() -> { + String traceId = MDC.get("AWS_LAMBDA_X_TRACE_ID"); + client.allTypes().join(); + }) + .join(); + + List requests = mockHttpClient.getRequests(); + + assertThat(requests).hasSize(2); + + assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); + assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); + + } finally { + MDC.clear(); + } + }); + } + + @Test + public void traceIdInterceptorPreservesTraceIdAcrossExceptionallyCompletedFutures() { + EnvironmentVariableHelper.run(env -> { + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-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("mdc-trace-123"); @@ -98,3 +191,4 @@ public void traceIdInterceptorPreservesTraceIdAcrossRetries() { }); } } + From 89883c7bd0758504c28ebbe06980d7beff98f8d0 Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Wed, 30 Jul 2025 21:22:04 -0700 Subject: [PATCH 5/6] Fix minor issues --- .../awscore/interceptor/TraceIdExecutionInterceptor.java | 5 ++++- .../java/software/amazon/awssdk/services/TraceIdTest.java | 4 +--- 2 files changed, 5 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 a16132a9d834..9fd6d211b735 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 @@ -70,7 +70,10 @@ public void onExecutionFailure(Context.FailedExecution context, ExecutionAttribu } private static void saveTraceId(ExecutionAttributes executionAttributes) { - MDC.put(CONCURRENT_TRACE_ID_KEY, executionAttributes.getAttribute(TRACE_ID)); + String traceId = executionAttributes.getAttribute(TRACE_ID); + if (traceId != null) { + MDC.put(CONCURRENT_TRACE_ID_KEY, executionAttributes.getAttribute(TRACE_ID)); + } } private Optional traceIdHeader(Context.ModifyHttpRequest context) { 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 13f799d63e0b..d01eaee17971 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 @@ -37,8 +37,7 @@ /** * Verifies that the {@link TraceIdExecutionInterceptor} is actually wired up for AWS services. */ -public class -TraceIdTest { +public class TraceIdTest { @Test public void traceIdInterceptorIsEnabled() { EnvironmentVariableHelper.run(env -> { @@ -130,7 +129,6 @@ public void traceIdInterceptorPreservesTraceIdAcrossChainedFutures() { client.allTypes() .thenRun(() -> { - String traceId = MDC.get("AWS_LAMBDA_X_TRACE_ID"); client.allTypes().join(); }) .join(); From 65fd82cbf41e1d34420079ebbbce11c4087a25cc Mon Sep 17 00:00:00 2001 From: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Date: Wed, 6 Aug 2025 15:18:30 -0700 Subject: [PATCH 6/6] Add test for pre-execution failures --- .../amazon/awssdk/services/TraceIdTest.java | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) 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 d01eaee17971..696d391d19fb 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,10 +18,14 @@ import static org.assertj.core.api.Assertions.assertThat; import java.util.List; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import org.slf4j.MDC; 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; @@ -188,5 +192,55 @@ public void traceIdInterceptorPreservesTraceIdAcrossExceptionallyCompletedFuture } }); } + + @Test + public void traceIdInterceptorPreservesTraceIdAcrossExceptionallyCompletedFuturesThrownInPreExecution() { + EnvironmentVariableHelper.run(env -> { + env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); + MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-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("mdc-trace-123"); + + } finally { + MDC.clear(); + } + }); + } }