Skip to content

Commit 65a5654

Browse files
RanVakninRanVaknin
authored
Feature/master/utils lite lambda trace (#6404)
* Add utils-lite package (#6395) Add utils-lite package * Add support for concurrent trace id propagation (#6403) Add support for concurrent trace id propagation * Add changelog * update utils-lite version to match parent * Fix again version * Bump version * Fix braziljson name * Add thread safe annotation * Adding a lazy loading of the hashmap * Refactored code according to comments * Modify threadLocal to InheritableThreadLocal * Add javadoc clarification * Add lambda check to prevent extra checks on non-lambda environments * Update beforeExecution logic * Fix checkstyle * Add test coverage * update version * Removing non-determinstic test * Removing unused dependencies * update version
1 parent cc6f80e commit 65a5654

File tree

15 files changed

+683
-6
lines changed

15 files changed

+683
-6
lines changed

‎.brazil.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
"sdk-core": { "packageName": "AwsJavaSdk-Core" },
3535
"url-connection-client": { "packageName": "AwsJavaSdk-HttpClient-UrlConnectionClient" },
3636
"utils": { "packageName": "AwsJavaSdk-Core-Utils" },
37+
"utils-lite": { "packageName": "AwsJavaSdk-Core-UtilsLite" },
3738
"imds": { "packageName": "AwsJavaSdk-Imds" },
3839
"crt-core": { "packageName": "AwsJavaSdk-Core-CrtCore" },
3940
"checksums-spi": { "packageName": "AwsJavaSdk-Core-ChecksumsSpi" },

‎.changes/next-release/feature-AWSSDKforJavav2-a3984eb.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "feature",
3+
"category": "AWS SDK for Java v2",
4+
"contributor": "",
5+
"description": "Add a utils-lite package that wraps threadlocal meant for internal shared usage across multiple applications across multiple applications"
6+
}

‎aws-sdk-java/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2113,6 +2113,11 @@ Amazon AutoScaling, etc).</description>
21132113
<artifactId>bcmdashboards</artifactId>
21142114
<version>${awsjavasdk.version}</version>
21152115
</dependency>
2116+
<dependency>
2117+
<groupId>software.amazon.awssdk</groupId>
2118+
<artifactId>utils-lite</artifactId>
2119+
<version>${awsjavasdk.version}</version>
2120+
</dependency>
21162121
</dependencies>
21172122
<build>
21182123
<finalName>${project.artifactId}-${project.version}</finalName>

‎bom/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,11 @@
212212
<artifactId>utils</artifactId>
213213
<version>${awsjavasdk.version}</version>
214214
</dependency>
215+
<dependency>
216+
<groupId>software.amazon.awssdk</groupId>
217+
<artifactId>utils-lite</artifactId>
218+
<version>${awsjavasdk.version}</version>
219+
</dependency>
215220
<dependency>
216221
<groupId>software.amazon.awssdk</groupId>
217222
<artifactId>cloudwatch-metric-publisher</artifactId>

‎core/aws-core/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,11 @@
113113
<groupId>software.amazon.eventstream</groupId>
114114
<artifactId>eventstream</artifactId>
115115
</dependency>
116+
<dependency>
117+
<groupId>software.amazon.awssdk</groupId>
118+
<artifactId>utils-lite</artifactId>
119+
<version>${awsjavasdk.version}</version>
120+
</dependency>
116121

117122
<dependency>
118123
<groupId>software.amazon.awssdk</groupId>

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

Lines changed: 47 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,40 +19,82 @@
1919
import software.amazon.awssdk.annotations.SdkProtectedApi;
2020
import software.amazon.awssdk.awscore.internal.interceptor.TracingSystemSetting;
2121
import software.amazon.awssdk.core.interceptor.Context;
22+
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
2223
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
2324
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
2425
import software.amazon.awssdk.http.SdkHttpRequest;
2526
import software.amazon.awssdk.utils.SystemSetting;
27+
import software.amazon.awssdk.utilslite.SdkInternalThreadLocal;
2628

2729
/**
2830
* The {@code TraceIdExecutionInterceptor} copies the trace details to the {@link #TRACE_ID_HEADER} header, assuming we seem to
29-
* be running in a lambda environment.
31+
* be running in a lambda environment.`
3032
*/
3133
@SdkProtectedApi
3234
public class TraceIdExecutionInterceptor implements ExecutionInterceptor {
3335
private static final String TRACE_ID_HEADER = "X-Amzn-Trace-Id";
3436
private static final String LAMBDA_FUNCTION_NAME_ENVIRONMENT_VARIABLE = "AWS_LAMBDA_FUNCTION_NAME";
37+
private static final String CONCURRENT_TRACE_ID_KEY = "AWS_LAMBDA_X_TRACE_ID";
38+
private static final ExecutionAttribute<String> TRACE_ID = new ExecutionAttribute<>("TraceId");
39+
40+
@Override
41+
public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) {
42+
if (lambdaFunctionNameEnvironmentVariable().isPresent()) {
43+
String traceId = SdkInternalThreadLocal.get(CONCURRENT_TRACE_ID_KEY);
44+
if (traceId != null) {
45+
executionAttributes.putAttribute(TRACE_ID, traceId);
46+
}
47+
}
48+
}
3549

3650
@Override
3751
public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
3852
Optional<String> traceIdHeader = traceIdHeader(context);
3953
if (!traceIdHeader.isPresent()) {
4054
Optional<String> lambdafunctionName = lambdaFunctionNameEnvironmentVariable();
41-
Optional<String> traceId = traceId();
55+
Optional<String> traceId = traceId(executionAttributes);
4256

4357
if (lambdafunctionName.isPresent() && traceId.isPresent()) {
4458
return context.httpRequest().copy(r -> r.putHeader(TRACE_ID_HEADER, traceId.get()));
4559
}
4660
}
47-
4861
return context.httpRequest();
4962
}
5063

64+
@Override
65+
public void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) {
66+
if (lambdaFunctionNameEnvironmentVariable().isPresent()) {
67+
saveTraceId(executionAttributes);
68+
}
69+
}
70+
71+
@Override
72+
public void onExecutionFailure(Context.FailedExecution context, ExecutionAttributes executionAttributes) {
73+
if (lambdaFunctionNameEnvironmentVariable().isPresent()) {
74+
saveTraceId(executionAttributes);
75+
}
76+
}
77+
78+
/**
79+
* Stores the trace ID in thread-local storage to ensure trace propagation across
80+
* thread boundaries during retries, or future chaining.
81+
*/
82+
private static void saveTraceId(ExecutionAttributes executionAttributes) {
83+
String traceId = executionAttributes.getAttribute(TRACE_ID);
84+
if (traceId != null) {
85+
SdkInternalThreadLocal.put(CONCURRENT_TRACE_ID_KEY, executionAttributes.getAttribute(TRACE_ID));
86+
}
87+
}
88+
5189
private Optional<String> traceIdHeader(Context.ModifyHttpRequest context) {
5290
return context.httpRequest().firstMatchingHeader(TRACE_ID_HEADER);
5391
}
5492

55-
private Optional<String> traceId() {
93+
private Optional<String> traceId(ExecutionAttributes executionAttributes) {
94+
Optional<String> traceId = Optional.ofNullable(executionAttributes.getAttribute(TRACE_ID));
95+
if (traceId.isPresent()) {
96+
return traceId;
97+
}
5698
return TracingSystemSetting._X_AMZN_TRACE_ID.getStringValue();
5799
}
58100

@@ -61,4 +103,4 @@ private Optional<String> lambdaFunctionNameEnvironmentVariable() {
61103
return SystemSetting.getStringValueFromEnvironmentVariable(LAMBDA_FUNCTION_NAME_ENVIRONMENT_VARIABLE);
62104
// CHECKSTYLE:ON
63105
}
64-
}
106+
}

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

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import software.amazon.awssdk.http.SdkHttpMethod;
2929
import software.amazon.awssdk.http.SdkHttpRequest;
3030
import software.amazon.awssdk.testutils.EnvironmentVariableHelper;
31+
import software.amazon.awssdk.utilslite.SdkInternalThreadLocal;
3132

3233
public class TraceIdExecutionInterceptorTest {
3334
@Test
@@ -111,6 +112,78 @@ public void headerNotAddedIfNoTraceIdEnvVar() {
111112
});
112113
}
113114

115+
@Test
116+
public void modifyHttpRequest_whenMultiConcurrencyModeWithInternalThreadLocal_shouldAddTraceIdHeader() {
117+
EnvironmentVariableHelper.run(env -> {
118+
resetRelevantEnvVars(env);
119+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
120+
SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123");
121+
122+
try {
123+
TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor();
124+
ExecutionAttributes executionAttributes = new ExecutionAttributes();
125+
126+
interceptor.beforeExecution(null, executionAttributes);
127+
Context.ModifyHttpRequest context = context();
128+
129+
SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes);
130+
assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123");
131+
} finally {
132+
SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID");
133+
}
134+
});
135+
}
136+
137+
@Test
138+
public void modifyHttpRequest_whenMultiConcurrencyModeWithBothInternalThreadLocalAndSystemProperty_shouldUseInternalThreadLocalValue() {
139+
EnvironmentVariableHelper.run(env -> {
140+
resetRelevantEnvVars(env);
141+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
142+
143+
SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "SdkInternalThreadLocal-trace-123");
144+
Properties props = System.getProperties();
145+
props.setProperty("com.amazonaws.xray.traceHeader", "sys-prop-345");
146+
147+
try {
148+
TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor();
149+
ExecutionAttributes executionAttributes = new ExecutionAttributes();
150+
151+
interceptor.beforeExecution(null, executionAttributes);
152+
153+
Context.ModifyHttpRequest context = context();
154+
SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes);
155+
156+
assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("SdkInternalThreadLocal-trace-123");
157+
} finally {
158+
SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID");
159+
props.remove("com.amazonaws.xray.traceHeader");
160+
}
161+
});
162+
}
163+
164+
@Test
165+
public void modifyHttpRequest_whenNotInLambdaEnvironmentWithInternalThreadLocal_shouldNotAddHeader() {
166+
EnvironmentVariableHelper.run(env -> {
167+
resetRelevantEnvVars(env);
168+
169+
SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "should-be-ignored");
170+
171+
try {
172+
TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor();
173+
ExecutionAttributes executionAttributes = new ExecutionAttributes();
174+
175+
interceptor.beforeExecution(null, executionAttributes);
176+
177+
Context.ModifyHttpRequest context = context();
178+
SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes);
179+
180+
assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).isEmpty();
181+
} finally {
182+
SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID");
183+
}
184+
});
185+
}
186+
114187
private Context.ModifyHttpRequest context() {
115188
return context(SdkHttpRequest.builder()
116189
.uri(URI.create("https://localhost"))

‎pom.xml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@
6767
<module>metric-publishers</module>
6868
<module>release-scripts</module>
6969
<module>utils</module>
70+
<module>utils-lite</module>
7071
<module>codegen-lite</module>
7172
<module>codegen-lite-maven-plugin</module>
7273
<module>archetypes</module>
@@ -664,6 +665,7 @@
664665
<includeModule>cloudwatch-metric-publisher</includeModule>
665666
<includeModule>emf-metric-logging-publisher</includeModule>
666667
<includeModule>utils</includeModule>
668+
<includeModule>utils-lite</includeModule>
667669
<includeModule>imds</includeModule>
668670
<includeModule>retries</includeModule>
669671
<includeModule>retries-spi</includeModule>

‎test/architecture-tests/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,11 @@
6161
<groupId>software.amazon.awssdk</groupId>
6262
<version>${awsjavasdk.version}</version>
6363
</dependency>
64+
<dependency>
65+
<artifactId>utils-lite</artifactId>
66+
<groupId>software.amazon.awssdk</groupId>
67+
<version>${awsjavasdk.version}</version>
68+
</dependency>
6469
<dependency>
6570
<artifactId>s3</artifactId>
6671
<groupId>software.amazon.awssdk</groupId>

‎test/architecture-tests/src/test/java/software/amazon/awssdk/archtests/UtilsLitePackageTest.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.archtests;
17+
18+
import static com.tngtech.archunit.lang.syntax.ArchRuleDefinition.classes;
19+
20+
import com.tngtech.archunit.core.domain.JavaClasses;
21+
import com.tngtech.archunit.core.importer.ClassFileImporter;
22+
import com.tngtech.archunit.lang.ArchRule;
23+
import org.junit.jupiter.api.Test;
24+
25+
/**
26+
* Architecture tests for the utils-lite package to ensure it only contains allowed classes.
27+
*/
28+
public class UtilsLitePackageTest {
29+
30+
private static final JavaClasses CLASSES = new ClassFileImporter()
31+
.importPackages("software.amazon.awssdk.utilslite");
32+
33+
@Test
34+
public void utilsLitePackage_shouldOnlyContainAllowedClasses() {
35+
ArchRule rule = classes()
36+
.that().resideInAPackage("software.amazon.awssdk.utilslite")
37+
.should().haveNameMatching(".*\\.(SdkInternalThreadLocal|SdkInternalThreadLocalTest)")
38+
.allowEmptyShould(true)
39+
.because("utils-lite package should only contain SdkInternalThreadLocal and its test");
40+
41+
rule.check(CLASSES);
42+
}
43+
}

0 commit comments

Comments
 (0)