Skip to content

Commit 7918d9a

Browse files
mricciamriccia
authored
fix: get trace id from system property when env var is not set (#1503)
* fix: check if XRAY Trace ID is present in System property * chore: remove erroneous extra char in tests
1 parent 0a30d8e commit 7918d9a

File tree

7 files changed

+77
-9
lines changed

7 files changed

+77
-9
lines changed

Diff for: ‎powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaConstants.java

+1
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ public class LambdaConstants {
2424
@Deprecated
2525
public static final String ON_DEMAND = "on-demand";
2626
public static final String X_AMZN_TRACE_ID = "_X_AMZN_TRACE_ID";
27+
public static final String XRAY_TRACE_HEADER = "com.amazonaws.xray.traceHeader";
2728
public static final String AWS_SAM_LOCAL = "AWS_SAM_LOCAL";
2829
public static final String ROOT_EQUALS = "Root=";
2930
public static final String POWERTOOLS_SERVICE_NAME = "POWERTOOLS_SERVICE_NAME";

Diff for: ‎powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessor.java

+7-1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
import static java.util.Optional.empty;
1818
import static java.util.Optional.of;
19+
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getProperty;
1920
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getenv;
2021

2122
import com.amazonaws.services.lambda.runtime.Context;
@@ -93,7 +94,12 @@ public static boolean isSamLocal() {
9394
}
9495

9596
public static Optional<String> getXrayTraceId() {
96-
final String X_AMZN_TRACE_ID = getenv(LambdaConstants.X_AMZN_TRACE_ID);
97+
String X_AMZN_TRACE_ID = getenv(LambdaConstants.X_AMZN_TRACE_ID);
98+
// For the Java Lambda 17+ runtime, the Trace ID is set as a System Property
99+
if (X_AMZN_TRACE_ID == null) {
100+
X_AMZN_TRACE_ID = getProperty(LambdaConstants.XRAY_TRACE_HEADER);
101+
}
102+
97103
if (X_AMZN_TRACE_ID != null) {
98104
return of(X_AMZN_TRACE_ID.split(";")[0].replace(LambdaConstants.ROOT_EQUALS, ""));
99105
}

Diff for: ‎powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/SystemWrapper.java

+4
Original file line numberDiff line numberDiff line change
@@ -21,4 +21,8 @@ private SystemWrapper() {
2121
public static String getenv(String name) {
2222
return System.getenv(name);
2323
}
24+
25+
public static String getProperty(String name) {
26+
return System.getProperty(name);
27+
}
2428
}

Diff for: ‎powertools-e2e-tests/src/test/java/software/amazon/lambda/powertools/LoggingE2ET.java

+2
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ public class LoggingE2ET {
4646
public static void setup() {
4747
infrastructure = Infrastructure.builder()
4848
.testName(LoggingE2ET.class.getSimpleName())
49+
.tracing(true)
4950
.pathToFunction("logging")
5051
.environmentVariables(
5152
Stream.of(new String[][] {
@@ -83,6 +84,7 @@ public void test_logInfoWithAdditionalKeys() throws JsonProcessingException {
8384
assertThat(jsonNode.get("message").asText()).isEqualTo("New Order");
8485
assertThat(jsonNode.get("orderId").asText()).isEqualTo(orderId);
8586
assertThat(jsonNode.get("coldStart").asBoolean()).isTrue();
87+
assertThat(jsonNode.get("xray_trace_id").asText()).isNotBlank();
8688
assertThat(jsonNode.get("function_request_id").asText()).isEqualTo(invocationResult1.getRequestId());
8789

8890
// second call should not be cold start

Diff for: ‎powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspectTest.java

+20-1
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import static org.mockito.Mockito.when;
2525
import static org.mockito.MockitoAnnotations.openMocks;
2626
import static org.skyscreamer.jsonassert.JSONAssert.assertEquals;
27+
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getProperty;
2728
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getenv;
2829

2930
import com.amazonaws.services.lambda.runtime.Context;
@@ -245,13 +246,31 @@ void shouldLogServiceNameWhenEnvVarSet() throws IllegalAccessException {
245246
.containsEntry("service", "testService");
246247
}
247248

249+
@Test
250+
void shouldLogxRayTraceIdSystemPropertySet() {
251+
String xRayTraceId = "1-5759e988-bd862e3fe1be46a994272793";
252+
253+
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class)) {
254+
mocked.when(() -> getenv("_X_AMZN_TRACE_ID"))
255+
.thenReturn(null);
256+
mocked.when(() -> getProperty("com.amazonaws.xray.traceHeader"))
257+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
258+
259+
requestHandler.handleRequest(new Object(), context);
260+
261+
assertThat(ThreadContext.getImmutableContext())
262+
.hasSize(EXPECTED_CONTEXT_SIZE + 1)
263+
.containsEntry("xray_trace_id", xRayTraceId);
264+
}
265+
}
266+
248267
@Test
249268
void shouldLogxRayTraceIdEnvVarSet() {
250269
String xRayTraceId = "1-5759e988-bd862e3fe1be46a994272793";
251270

252271
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class)) {
253272
mocked.when(() -> getenv("_X_AMZN_TRACE_ID"))
254-
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
273+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
255274

256275
requestHandler.handleRequest(new Object(), context);
257276

Diff for: ‎powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/MetricsLoggerTest.java

+40-4
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import static org.assertj.core.api.Assertions.assertThat;
1919
import static org.assertj.core.api.Assertions.assertThatNullPointerException;
2020
import static org.mockito.Mockito.mockStatic;
21+
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getProperty;
2122
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getenv;
2223

2324
import com.fasterxml.jackson.core.JsonProcessingException;
@@ -66,7 +67,7 @@ void singleMetricsCaptureUtilityWithDefaultDimension() {
6667
software.amazon.lambda.powertools.core.internal.SystemWrapper.class)) {
6768
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
6869
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
69-
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
70+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
7071

7172
MetricsUtils.defaultDimensions(DimensionSet.of("Service", "Booking"));
7273

@@ -96,7 +97,7 @@ void singleMetricsCaptureUtility() {
9697
software.amazon.lambda.powertools.core.internal.SystemWrapper.class)) {
9798
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
9899
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
99-
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
100+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
100101

101102
MetricsUtils.withSingleMetric("Metric1", 1, Unit.COUNT, "test",
102103
metricsLogger -> metricsLogger.setDimensions(DimensionSet.of("Dimension1", "Value1")));
@@ -123,7 +124,7 @@ void singleMetricsCaptureUtilityWithDefaultNameSpace() {
123124
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
124125
mocked.when(() -> SystemWrapper.getenv("POWERTOOLS_METRICS_NAMESPACE")).thenReturn("GlobalName");
125126
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
126-
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
127+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
127128

128129
MetricsUtils.withSingleMetric("Metric1", 1, Unit.COUNT,
129130
metricsLogger -> metricsLogger.setDimensions(DimensionSet.of("Dimension1", "Value1")));
@@ -165,14 +166,49 @@ void shouldThrowExceptionWhenDefaultDimensionIsNull() {
165166
.withMessage("Null dimension set not allowed");
166167
}
167168

169+
@Test
170+
void shouldUseTraceIdFromSystemPropertyIfEnvVarNotPresent() {
171+
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class);
172+
MockedStatic<software.amazon.lambda.powertools.core.internal.SystemWrapper> internalWrapper = mockStatic(
173+
software.amazon.lambda.powertools.core.internal.SystemWrapper.class)) {
174+
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
175+
mocked.when(() -> SystemWrapper.getenv("POWERTOOLS_METRICS_NAMESPACE")).thenReturn("GlobalName");
176+
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
177+
.thenReturn(null);
178+
internalWrapper.when(() -> getProperty("com.amazonaws.xray.traceHeader"))
179+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
180+
181+
MetricsUtils.withSingleMetric("Metric1", 1, Unit.COUNT,
182+
metricsLogger -> metricsLogger.setDimensions(DimensionSet.of("Dimension1", "Value1")));
183+
184+
assertThat(out.toString())
185+
.satisfies(s ->
186+
{
187+
Map<String, Object> logAsJson = readAsJson(s);
188+
189+
assertThat(logAsJson)
190+
.containsEntry("Metric1", 1.0)
191+
.containsEntry("Dimension1", "Value1")
192+
.containsKey("_aws")
193+
.containsEntry("xray_trace_id", "1-5759e988-bd862e3fe1be46a994272793");
194+
195+
Map<String, Object> aws = (Map<String, Object>) logAsJson.get("_aws");
196+
197+
assertThat(aws.get("CloudWatchMetrics"))
198+
.asString()
199+
.contains("Namespace=GlobalName");
200+
});
201+
}
202+
}
203+
168204
private void testLogger(Consumer<Consumer<MetricsLogger>> methodToTest) {
169205
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class);
170206
MockedStatic<software.amazon.lambda.powertools.core.internal.SystemWrapper> internalWrapper = mockStatic(
171207
software.amazon.lambda.powertools.core.internal.SystemWrapper.class)) {
172208
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
173209
mocked.when(() -> SystemWrapper.getenv("POWERTOOLS_METRICS_NAMESPACE")).thenReturn("GlobalName");
174210
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
175-
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
211+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
176212

177213
methodToTest.accept(metricsLogger ->
178214
{

Diff for: ‎powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspectTest.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ public void metricsWithoutColdStart() {
9191

9292
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
9393
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
94-
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
94+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
9595

9696
MetricsUtils.defaultDimensions(null);
9797
requestHandler = new PowertoolsMetricsEnabledHandler();
@@ -135,7 +135,7 @@ public void metricsWithDefaultDimensionSpecified() {
135135

136136
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
137137
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
138-
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
138+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
139139

140140
requestHandler = new PowertoolsMetricsEnabledDefaultDimensionHandler();
141141

@@ -179,7 +179,7 @@ public void metricsWithDefaultNoDimensionSpecified() {
179179

180180
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
181181
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
182-
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
182+
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
183183

184184
requestHandler = new PowertoolsMetricsEnabledDefaultNoDimensionHandler();
185185

0 commit comments

Comments
 (0)