Skip to content

Commit 55038f7

Browse files
jj22eejj22ee
authored
Support extracting X-Ray Trace Id via SdkInternalThreadLocal in Lambda Segment Context (#433)
1 parent 0ae0b07 commit 55038f7

File tree

3 files changed

+129
-5
lines changed

3 files changed

+129
-5
lines changed

‎aws-xray-recorder-sdk-core/build.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ plugins {
66
dependencies {
77
api("commons-logging:commons-logging:1.3.5")
88

9+
implementation("software.amazon.awssdk:utils-lite:2.34.0")
910
implementation("com.fasterxml.jackson.core:jackson-annotations:2.17.0")
1011
implementation("com.fasterxml.jackson.core:jackson-databind:2.17.0")
1112
implementation("com.google.auto.value:auto-value-annotations:1.10.4")

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,20 +30,28 @@
3030
import java.util.Objects;
3131
import org.apache.commons.logging.Log;
3232
import org.apache.commons.logging.LogFactory;
33+
import software.amazon.awssdk.utilslite.SdkInternalThreadLocal;
3334

3435
public class LambdaSegmentContext implements SegmentContext {
3536
private static final Log logger = LogFactory.getLog(LambdaSegmentContext.class);
3637

3738
private static final String LAMBDA_TRACE_HEADER_KEY = "_X_AMZN_TRACE_ID";
38-
39+
private static final String CONCURRENT_TRACE_ID_KEY = "AWS_LAMBDA_X_TRACE_ID";
40+
3941
// See: https://github.com/aws/aws-xray-sdk-java/issues/251
4042
private static final String LAMBDA_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader";
4143

4244
public static TraceHeader getTraceHeaderFromEnvironment() {
43-
String lambdaTraceHeaderKey = System.getenv(LAMBDA_TRACE_HEADER_KEY);
44-
return TraceHeader.fromString(lambdaTraceHeaderKey != null && lambdaTraceHeaderKey.length() > 0
45-
? lambdaTraceHeaderKey
46-
: System.getProperty(LAMBDA_TRACE_HEADER_PROP));
45+
String lambdaTraceHeaderKeyFromMdc = SdkInternalThreadLocal.get(CONCURRENT_TRACE_ID_KEY);
46+
String lambdaTraceHeaderKeyFromEnvVar = System.getenv(LAMBDA_TRACE_HEADER_KEY);
47+
48+
if (lambdaTraceHeaderKeyFromMdc != null && lambdaTraceHeaderKeyFromMdc.length() > 0) {
49+
return TraceHeader.fromString(lambdaTraceHeaderKeyFromMdc);
50+
} else if (lambdaTraceHeaderKeyFromEnvVar != null && lambdaTraceHeaderKeyFromEnvVar.length() > 0) {
51+
return TraceHeader.fromString(lambdaTraceHeaderKeyFromEnvVar);
52+
} else {
53+
return TraceHeader.fromString(System.getProperty(LAMBDA_TRACE_HEADER_PROP));
54+
}
4755
}
4856

4957
// SuppressWarnings is needed for passing Root TraceId to noOp segment

‎aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import org.mockito.junit.jupiter.MockitoExtension;
4141
import org.mockito.junit.jupiter.MockitoSettings;
4242
import org.mockito.quality.Strictness;
43+
import software.amazon.awssdk.utilslite.SdkInternalThreadLocal;
4344

4445
@ExtendWith(MockitoExtension.class)
4546
@MockitoSettings(strictness = Strictness.LENIENT)
@@ -174,4 +175,118 @@ private static void testContextResultsInNoOpSegmentParent() {
174175
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
175176
assertThat(AWSXRay.getTraceEntity()).isNull();
176177
}
178+
179+
@Test
180+
@SetSystemProperty(key = "com.amazonaws.xray.traceHeader", value = TRACE_HEADER)
181+
void testSystemPropertyFallbackWithTraceValidation() {
182+
LambdaSegmentContext mockContext = new LambdaSegmentContext();
183+
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
184+
FacadeSegment parent = (FacadeSegment) subsegment.getParent();
185+
186+
// Verify system property values are used correctly
187+
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
188+
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
189+
assertThat(parent.isSampled()).isTrue();
190+
191+
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
192+
}
193+
194+
@Test
195+
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER)
196+
void testEnvironmentVariableFallbackWithTraceValidation() {
197+
LambdaSegmentContext mockContext = new LambdaSegmentContext();
198+
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
199+
FacadeSegment parent = (FacadeSegment) subsegment.getParent();
200+
201+
// Verify system property values are used correctly
202+
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
203+
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
204+
assertThat(parent.isSampled()).isTrue();
205+
206+
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
207+
}
208+
209+
@Test
210+
@SetSystemProperty(key = "com.amazonaws.xray.traceHeader", value = TRACE_HEADER_2)
211+
void testSdkInternalThreadLocalTakesPriorityOverSystemProperty() {
212+
SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", TRACE_HEADER);
213+
214+
try {
215+
LambdaSegmentContext mockContext = new LambdaSegmentContext();
216+
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
217+
FacadeSegment parent = (FacadeSegment) subsegment.getParent();
218+
219+
// Verify SdkInternalThreadLocal values are used (TRACE_HEADER), not system property values (TRACE_HEADER_2)
220+
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
221+
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
222+
assertThat(parent.isSampled()).isTrue();
223+
224+
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
225+
} finally {
226+
SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID");
227+
}
228+
}
229+
230+
@Test
231+
@SetSystemProperty(key = "com.amazonaws.xray.traceHeader", value = TRACE_HEADER)
232+
void testSdkInternalThreadLocalWithEmptyStringFallsBackToSystemProperty() {
233+
SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "");
234+
235+
try {
236+
LambdaSegmentContext mockContext = new LambdaSegmentContext();
237+
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
238+
FacadeSegment parent = (FacadeSegment) subsegment.getParent();
239+
240+
// Verify system property values are used as fallback when SdkInternalThreadLocal returns empty string
241+
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
242+
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
243+
assertThat(parent.isSampled()).isTrue();
244+
245+
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
246+
} finally {
247+
SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID");
248+
}
249+
}
250+
251+
@Test
252+
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER_2)
253+
void testSdkInternalThreadLocalTakesPriorityOverEnvironmentVariable() {
254+
SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", TRACE_HEADER);
255+
256+
try {
257+
LambdaSegmentContext mockContext = new LambdaSegmentContext();
258+
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
259+
FacadeSegment parent = (FacadeSegment) subsegment.getParent();
260+
261+
// Verify SdkInternalThreadLocal values are used (TRACE_HEADER), not environment variable values (TRACE_HEADER_2)
262+
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
263+
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
264+
assertThat(parent.isSampled()).isTrue();
265+
266+
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
267+
} finally {
268+
SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID");
269+
}
270+
}
271+
272+
@Test
273+
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER)
274+
void testSdkInternalThreadLocalWithEmptyStringFallsBackToEnvironmentVariable() {
275+
SdkInternalThreadLocal.put("AWS_LAMBDA_X_TRACE_ID", "");
276+
277+
try {
278+
LambdaSegmentContext mockContext = new LambdaSegmentContext();
279+
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
280+
FacadeSegment parent = (FacadeSegment) subsegment.getParent();
281+
282+
// Verify environment variable values are used as fallback when SdkInternalThreadLocal returns empty string
283+
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
284+
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
285+
assertThat(parent.isSampled()).isTrue();
286+
287+
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
288+
} finally {
289+
SdkInternalThreadLocal.remove("AWS_LAMBDA_X_TRACE_ID");
290+
}
291+
}
177292
}

0 commit comments

Comments
 (0)