Skip to content

Feature/master/utils lite lambda trace again #6433

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
joviegas merged 22 commits into from
Sep 19, 2025
Merged

Feature/master/utils lite lambda trace again #6433

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
8244f6b
Add utils-lite package (#6395)
RanVaknin Sep 5, 2025
eb69c58
Add support for concurrent trace id propagation (#6403)
RanVaknin Sep 5, 2025
f9e3ec6
Add changelog
RanVaknin Sep 5, 2025
c6ea4cc
update utils-lite version to match parent
RanVaknin Sep 5, 2025
f505c5c
Fix again version
RanVaknin Sep 5, 2025
c8b62a0
Bump version
RanVaknin Sep 11, 2025
fc8ddd0
Fix braziljson name
RanVaknin Sep 11, 2025
010f789
Add thread safe annotation
RanVaknin Sep 11, 2025
57976e7
Adding a lazy loading of the hashmap
RanVaknin Sep 11, 2025
41684b7
Refactored code according to comments
RanVaknin Sep 12, 2025
0da1f27
Modify threadLocal to InheritableThreadLocal
RanVaknin Sep 12, 2025
42e6110
Add javadoc clarification
RanVaknin Sep 15, 2025
514fce8
Add lambda check to prevent extra checks on non-lambda environments
RanVaknin Sep 15, 2025
14bb60c
Update beforeExecution logic
RanVaknin Sep 15, 2025
9d1a753
Fix checkstyle
RanVaknin Sep 15, 2025
842e843
Add test coverage
RanVaknin Sep 17, 2025
197d09b
update version
RanVaknin Sep 17, 2025
fd42b97
Removing non-determinstic test
RanVaknin Sep 17, 2025
8a34b60
Removing unused dependencies
RanVaknin Sep 17, 2025
e7c081b
update version
RanVaknin Sep 18, 2025
61e037e
Add support for lambda traces via SdkInternalThreadLocal
RanVaknin Sep 19, 2025
f02ff03
Bump version
RanVaknin Sep 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .brazil.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
"sdk-core": { "packageName": "AwsJavaSdk-Core" },
"url-connection-client": { "packageName": "AwsJavaSdk-HttpClient-UrlConnectionClient" },
"utils": { "packageName": "AwsJavaSdk-Core-Utils" },
"utils-lite": { "packageName": "AwsJavaSdk-Core-UtilsLite" },
"imds": { "packageName": "AwsJavaSdk-Imds" },
"crt-core": { "packageName": "AwsJavaSdk-Core-CrtCore" },
"checksums-spi": { "packageName": "AwsJavaSdk-Core-ChecksumsSpi" },
Expand Down
6 changes: 6 additions & 0 deletions .changes/next-release/feature-AWSSDKforJavav2-a3984eb.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
{
"type": "feature",
"category": "AWS SDK for Java v2",
"contributor": "",
"description": "Add a utils-lite package that wraps threadlocal meant for internal shared usage across multiple applications across multiple applications"
}
5 changes: 5 additions & 0 deletions aws-sdk-java/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2113,6 +2113,11 @@ Amazon AutoScaling, etc).</description>
<artifactId>bcmdashboards</artifactId>
<version>${awsjavasdk.version}</version>
</dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>utils-lite</artifactId>
<version>${awsjavasdk.version}</version>
</dependency>
</dependencies>
<build>
<finalName>${project.artifactId}-${project.version}</finalName>
Expand Down
5 changes: 5 additions & 0 deletions bom/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,11 @@
<artifactId>utils</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>cloudwatch-metric-publisher</artifactId>
Expand Down
5 changes: 5 additions & 0 deletions core/aws-core/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +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>
Expand Down
52 changes: 47 additions & 5 deletions ...src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,40 +19,82 @@
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
* be running in a lambda environment.
* be running in a lambda environment.`
*/
@SdkProtectedApi
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) {
if (lambdaFunctionNameEnvironmentVariable().isPresent()) {
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) {
if (lambdaFunctionNameEnvironmentVariable().isPresent()) {
saveTraceId(executionAttributes);
}
}

@Override
public void onExecutionFailure(Context.FailedExecution context, ExecutionAttributes executionAttributes) {
if (lambdaFunctionNameEnvironmentVariable().isPresent()) {
saveTraceId(executionAttributes);
}
}

/**
* Stores the trace ID in thread-local storage to ensure trace propagation across
* thread boundaries during retries, or future chaining.
*/
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();
}

Expand All @@ -61,4 +103,4 @@ private Optional<String> lambdaFunctionNameEnvironmentVariable() {
return SystemSetting.getStringValueFromEnvironmentVariable(LAMBDA_FUNCTION_NAME_ENVIRONMENT_VARIABLE);
// CHECKSTYLE:ON
}
}
}
73 changes: 73 additions & 0 deletions ...test/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"))
Expand Down
2 changes: 2 additions & 0 deletions pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -664,6 +665,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>
Expand Down
5 changes: 5 additions & 0 deletions test/architecture-tests/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
43 changes: 43 additions & 0 deletions ...chitecture-tests/src/test/java/software/amazon/awssdk/archtests/UtilsLitePackageTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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);
}
}
Loading
Loading