Skip to content

Commit 71aacb4

Browse files
phipagkjswaruph
phipag
authored and
kjswaruph
committed
improv: Make trace id access more robust. (aws-powertools#2188)
* improv: Make trace id access more robust. * Update comments.
1 parent d2e5023 commit 71aacb4

File tree

20 files changed

+91
-45
lines changed

20 files changed

+91
-45
lines changed

‎examples/powertools-examples-serialization/sam-graalvm/src/main/resources/META-INF/native-image/com.amazonaws/aws-lambda-java-runtime-interface-client/reflect-config.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,6 @@
3333
},
3434
{
3535
"name":"software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor",
36-
"fields":[{"name":"IS_COLD_START"},{"name":"SERVICE_NAME"}]
36+
"fields":[{"name":"isColdStart"},{"name":"serviceName"}]
3737
}
38-
]
38+
]

‎powertools-common/pom.xml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,10 @@
4444
<groupId>com.amazonaws</groupId>
4545
<artifactId>aws-lambda-java-core</artifactId>
4646
</dependency>
47+
<dependency>
48+
<groupId>software.amazon.awssdk</groupId>
49+
<artifactId>utils-lite</artifactId>
50+
</dependency>
4751

4852
<!-- Test dependencies -->
4953
<dependency>

‎powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaConstants.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,16 @@
1414

1515
package software.amazon.lambda.powertools.common.internal;
1616

17-
public class LambdaConstants {
17+
public final class LambdaConstants {
18+
private LambdaConstants() {
19+
// Constant holder class
20+
}
21+
1822
public static final String LAMBDA_FUNCTION_NAME_ENV = "AWS_LAMBDA_FUNCTION_NAME";
1923
public static final String AWS_REGION_ENV = "AWS_REGION";
2024
public static final String X_AMZN_TRACE_ID = "_X_AMZN_TRACE_ID";
2125
public static final String XRAY_TRACE_HEADER = "com.amazonaws.xray.traceHeader";
26+
public static final String AWS_LAMBDA_X_TRACE_ID = "AWS_LAMBDA_X_TRACE_ID";
2227
public static final String AWS_SAM_LOCAL = "AWS_SAM_LOCAL";
2328
public static final String ROOT_EQUALS = "Root=";
2429
public static final String POWERTOOLS_SERVICE_NAME = "POWERTOOLS_SERVICE_NAME";

‎powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessor.java‎

Lines changed: 27 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -19,28 +19,33 @@
1919
import static software.amazon.lambda.powertools.common.internal.SystemWrapper.getProperty;
2020
import static software.amazon.lambda.powertools.common.internal.SystemWrapper.getenv;
2121

22-
import com.amazonaws.services.lambda.runtime.Context;
23-
import com.amazonaws.services.lambda.runtime.RequestHandler;
24-
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
2522
import java.io.InputStream;
2623
import java.io.OutputStream;
2724
import java.util.Optional;
25+
2826
import org.aspectj.lang.ProceedingJoinPoint;
2927

28+
import com.amazonaws.services.lambda.runtime.Context;
29+
import com.amazonaws.services.lambda.runtime.RequestHandler;
30+
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
31+
32+
import software.amazon.awssdk.utilslite.SdkInternalThreadLocal;
33+
3034
public final class LambdaHandlerProcessor {
3135

32-
// SERVICE_NAME cannot be final for testing purposes
33-
private static String SERVICE_NAME = calculateServiceName();
36+
// serviceName cannot be final for testing purposes
37+
private static String serviceName = calculateServiceName();
3438

35-
private static Boolean IS_COLD_START = null;
39+
private static Boolean isColdStart = null;
3640

3741
private LambdaHandlerProcessor() {
3842
// Hide default constructor
3943
}
4044

4145
private static String calculateServiceName() {
4246
return null != getenv(LambdaConstants.POWERTOOLS_SERVICE_NAME)
43-
? getenv(LambdaConstants.POWERTOOLS_SERVICE_NAME) : LambdaConstants.SERVICE_UNDEFINED;
47+
? getenv(LambdaConstants.POWERTOOLS_SERVICE_NAME)
48+
: LambdaConstants.SERVICE_UNDEFINED;
4449
}
4550

4651
public static boolean isHandlerMethod(final ProceedingJoinPoint pjp) {
@@ -79,35 +84,41 @@ public static Context extractContext(final ProceedingJoinPoint pjp) {
7984
}
8085

8186
public static String serviceName() {
82-
return SERVICE_NAME;
87+
return serviceName;
8388
}
8489

8590
// Method used for testing purposes
8691
protected static void resetServiceName() {
87-
SERVICE_NAME = calculateServiceName();
92+
serviceName = calculateServiceName();
8893
}
8994

9095
public static boolean isColdStart() {
91-
return IS_COLD_START == null;
96+
return isColdStart == null;
9297
}
9398

9499
public static void coldStartDone() {
95-
IS_COLD_START = false;
100+
isColdStart = false;
96101
}
97102

98103
public static boolean isSamLocal() {
99104
return "true".equals(getenv(LambdaConstants.AWS_SAM_LOCAL));
100105
}
101106

102107
public static Optional<String> getXrayTraceId() {
103-
String X_AMZN_TRACE_ID = getenv(LambdaConstants.X_AMZN_TRACE_ID);
108+
// Try SdkInternalThreadLocal first
109+
String traceId = SdkInternalThreadLocal.get(LambdaConstants.AWS_LAMBDA_X_TRACE_ID);
110+
111+
// Fallback to environment based approach
112+
if (traceId == null) {
113+
traceId = getenv(LambdaConstants.X_AMZN_TRACE_ID);
114+
}
104115
// For the Java Lambda 17+ runtime, the Trace ID is set as a System Property
105-
if (X_AMZN_TRACE_ID == null) {
106-
X_AMZN_TRACE_ID = getProperty(LambdaConstants.XRAY_TRACE_HEADER);
116+
if (traceId == null) {
117+
traceId = getProperty(LambdaConstants.XRAY_TRACE_HEADER);
107118
}
108119

109-
if (X_AMZN_TRACE_ID != null) {
110-
return of(X_AMZN_TRACE_ID.split(";")[0].replace(LambdaConstants.ROOT_EQUALS, ""));
120+
if (traceId != null) {
121+
return of(traceId.split(";")[0].replace(LambdaConstants.ROOT_EQUALS, ""));
111122
}
112123
return empty();
113124
}

‎powertools-common/src/test/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessorTest.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,18 +25,26 @@
2525

2626
import org.aspectj.lang.ProceedingJoinPoint;
2727
import org.aspectj.lang.Signature;
28+
import org.junit.jupiter.api.AfterEach;
2829
import org.junit.jupiter.api.Test;
2930
import org.junitpioneer.jupiter.ClearEnvironmentVariable;
31+
import org.junitpioneer.jupiter.ClearSystemProperty;
3032
import org.junitpioneer.jupiter.SetEnvironmentVariable;
3133

3234
import com.amazonaws.services.lambda.runtime.Context;
3335
import com.amazonaws.services.lambda.runtime.RequestHandler;
3436
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
3537

38+
import software.amazon.awssdk.utilslite.SdkInternalThreadLocal;
3639
import software.amazon.lambda.powertools.common.stubs.TestLambdaContext;
3740

3841
class LambdaHandlerProcessorTest {
3942

43+
@AfterEach
44+
void cleanup() {
45+
SdkInternalThreadLocal.clear();
46+
}
47+
4048
@Test
4149
void isHandlerMethod_shouldRecognizeRequestHandler() {
4250
Context context = new TestLambdaContext();
@@ -159,6 +167,24 @@ void getXrayTraceId_notPresent() {
159167
assertThat(isXRayTraceIdPresent).isFalse();
160168
}
161169

170+
@Test
171+
@ClearEnvironmentVariable(key = LambdaConstants.X_AMZN_TRACE_ID)
172+
@ClearSystemProperty(key = LambdaConstants.XRAY_TRACE_HEADER)
173+
void getXrayTraceId_fromSdkInternalThreadLocal() {
174+
// Verify no trace ID initially
175+
assertThat(LambdaHandlerProcessor.getXrayTraceId()).isEmpty();
176+
177+
// Set trace ID in SdkInternalThreadLocal
178+
String expectedTraceId = "1-5759e988-bd862e3fe1be46a994272793";
179+
SdkInternalThreadLocal.put(LambdaConstants.AWS_LAMBDA_X_TRACE_ID,
180+
"Root=" + expectedTraceId + ";Parent=53995c3f42cd8ad8;Sampled=1");
181+
182+
// Verify trace ID is now present
183+
Optional<String> traceId = LambdaHandlerProcessor.getXrayTraceId();
184+
assertThat(traceId).isPresent();
185+
assertThat(traceId.get()).isEqualTo(expectedTraceId);
186+
}
187+
162188
@Test
163189
void extractContext_fromRequestHandler() {
164190
Object[] args = { new Object(), new TestLambdaContext() };

‎powertools-logging/powertools-logging-log4j/src/main/resources/META-INF/native-image/software.amazon.lambda/powertools-logging-log4j/reflect-config.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1077,7 +1077,7 @@
10771077
},
10781078
{
10791079
"name":"software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor",
1080-
"fields":[{"name":"IS_COLD_START"}]
1080+
"fields":[{"name":"isColdStart"}]
10811081
},
10821082
{
10831083
"name":"software.amazon.lambda.powertools.logging.log4j.BufferingAppender",

‎powertools-logging/powertools-logging-log4j/src/test/java/org/apache/logging/log4j/layout/template/json/resolver/PowerToolsResolverFactoryTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ class PowerToolsResolverFactoryTest {
4545
@BeforeEach
4646
void setUp() throws IllegalAccessException, IOException {
4747
MDC.clear();
48-
writeStaticField(LambdaHandlerProcessor.class, "IS_COLD_START", null, true);
48+
writeStaticField(LambdaHandlerProcessor.class, "isColdStart", null, true);
4949
context = new TestLambdaContext();
5050
// Make sure file is cleaned up before running tests
5151
try {

‎powertools-logging/powertools-logging-logback/src/main/resources/META-INF/native-image/software.amazon.lambda/powertools-logging-logback/reflect-config.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -163,7 +163,7 @@
163163
},
164164
{
165165
"name":"software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor",
166-
"fields":[{"name":"IS_COLD_START"}]
166+
"fields":[{"name":"isColdStart"}]
167167
},
168168
{
169169
"name":"software.amazon.lambda.powertools.logging.logback.BufferingAppender",

‎powertools-logging/powertools-logging-logback/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaEcsEncoderTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ class LambdaEcsEncoderTest {
5151
@BeforeEach
5252
void setUp() throws IllegalAccessException, IOException {
5353
MDC.clear();
54-
writeStaticField(LambdaHandlerProcessor.class, "IS_COLD_START", null, true);
54+
writeStaticField(LambdaHandlerProcessor.class, "isColdStart", null, true);
5555
context = new TestLambdaContext();
5656
// Make sure file is cleaned up before running tests
5757
try {

‎powertools-logging/powertools-logging-logback/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaJsonEncoderTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ class LambdaJsonEncoderTest {
7373
@BeforeEach
7474
void setUp() throws IllegalAccessException, IOException {
7575
MDC.clear();
76-
writeStaticField(LambdaHandlerProcessor.class, "IS_COLD_START", null, true);
76+
writeStaticField(LambdaHandlerProcessor.class, "isColdStart", null, true);
7777
context = new TestLambdaContext();
7878
// Make sure file is cleaned up before running tests
7979
try {

0 commit comments

Comments
 (0)