Skip to content

Commit

Permalink
fix: get trace id from system property when env var is not set (#1503)
Browse files Browse the repository at this point in the history 
* fix: check if XRAY Trace ID is present in System property

* chore: remove erroneous extra char in tests
  • Loading branch information
@mriccia
mriccia committed Nov 8, 2023
1 parent 0a30d8e commit 7918d9a
Show file tree
Hide file tree
Showing 7 changed files with 77 additions and 9 deletions.
1 change: 1 addition & 0 deletions ...s-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaConstants.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ public class LambdaConstants {
@Deprecated
public static final String ON_DEMAND = "on-demand";
public static final String X_AMZN_TRACE_ID = "_X_AMZN_TRACE_ID";
public static final String XRAY_TRACE_HEADER = "com.amazonaws.xray.traceHeader";
public static final String AWS_SAM_LOCAL = "AWS_SAM_LOCAL";
public static final String ROOT_EQUALS = "Root=";
public static final String POWERTOOLS_SERVICE_NAME = "POWERTOOLS_SERVICE_NAME";
Expand Down
8 changes: 7 additions & 1 deletion ...src/main/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import static java.util.Optional.empty;
import static java.util.Optional.of;
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getProperty;
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getenv;

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

public static Optional<String> getXrayTraceId() {
final String X_AMZN_TRACE_ID = getenv(LambdaConstants.X_AMZN_TRACE_ID);
String X_AMZN_TRACE_ID = getenv(LambdaConstants.X_AMZN_TRACE_ID);
// For the Java Lambda 17+ runtime, the Trace ID is set as a System Property
if (X_AMZN_TRACE_ID == null) {
X_AMZN_TRACE_ID = getProperty(LambdaConstants.XRAY_TRACE_HEADER);
}

if (X_AMZN_TRACE_ID != null) {
return of(X_AMZN_TRACE_ID.split(";")[0].replace(LambdaConstants.ROOT_EQUALS, ""));
}
Expand Down
4 changes: 4 additions & 0 deletions ...ols-core/src/main/java/software/amazon/lambda/powertools/core/internal/SystemWrapper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,4 +21,8 @@ private SystemWrapper() {
public static String getenv(String name) {
return System.getenv(name);
}

public static String getProperty(String name) {
return System.getProperty(name);
}
}
2 changes: 2 additions & 0 deletions powertools-e2e-tests/src/test/java/software/amazon/lambda/powertools/LoggingE2ET.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ public class LoggingE2ET {
public static void setup() {
infrastructure = Infrastructure.builder()
.testName(LoggingE2ET.class.getSimpleName())
.tracing(true)
.pathToFunction("logging")
.environmentVariables(
Stream.of(new String[][] {
Expand Down Expand Up @@ -83,6 +84,7 @@ public void test_logInfoWithAdditionalKeys() throws JsonProcessingException {
assertThat(jsonNode.get("message").asText()).isEqualTo("New Order");
assertThat(jsonNode.get("orderId").asText()).isEqualTo(orderId);
assertThat(jsonNode.get("coldStart").asBoolean()).isTrue();
assertThat(jsonNode.get("xray_trace_id").asText()).isNotBlank();
assertThat(jsonNode.get("function_request_id").asText()).isEqualTo(invocationResult1.getRequestId());

// second call should not be cold start
Expand Down
21 changes: 20 additions & 1 deletion ...test/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspectTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import static org.mockito.Mockito.when;
import static org.mockito.MockitoAnnotations.openMocks;
import static org.skyscreamer.jsonassert.JSONAssert.assertEquals;
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getProperty;
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getenv;

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

@Test
void shouldLogxRayTraceIdSystemPropertySet() {
String xRayTraceId = "1-5759e988-bd862e3fe1be46a994272793";

try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class)) {
mocked.when(() -> getenv("_X_AMZN_TRACE_ID"))
.thenReturn(null);
mocked.when(() -> getProperty("com.amazonaws.xray.traceHeader"))
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");

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

assertThat(ThreadContext.getImmutableContext())
.hasSize(EXPECTED_CONTEXT_SIZE + 1)
.containsEntry("xray_trace_id", xRayTraceId);
}
}

@Test
void shouldLogxRayTraceIdEnvVarSet() {
String xRayTraceId = "1-5759e988-bd862e3fe1be46a994272793";

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

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

Expand Down
44 changes: 40 additions & 4 deletions ...ls-metrics/src/test/java/software/amazon/lambda/powertools/metrics/MetricsLoggerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatNullPointerException;
import static org.mockito.Mockito.mockStatic;
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getProperty;
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getenv;

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

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

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

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

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

@Test
void shouldUseTraceIdFromSystemPropertyIfEnvVarNotPresent() {
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class);
MockedStatic<software.amazon.lambda.powertools.core.internal.SystemWrapper> internalWrapper = mockStatic(
software.amazon.lambda.powertools.core.internal.SystemWrapper.class)) {
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
mocked.when(() -> SystemWrapper.getenv("POWERTOOLS_METRICS_NAMESPACE")).thenReturn("GlobalName");
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
.thenReturn(null);
internalWrapper.when(() -> getProperty("com.amazonaws.xray.traceHeader"))
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");

MetricsUtils.withSingleMetric("Metric1", 1, Unit.COUNT,
metricsLogger -> metricsLogger.setDimensions(DimensionSet.of("Dimension1", "Value1")));

assertThat(out.toString())
.satisfies(s ->
{
Map<String, Object> logAsJson = readAsJson(s);

assertThat(logAsJson)
.containsEntry("Metric1", 1.0)
.containsEntry("Dimension1", "Value1")
.containsKey("_aws")
.containsEntry("xray_trace_id", "1-5759e988-bd862e3fe1be46a994272793");

Map<String, Object> aws = (Map<String, Object>) logAsJson.get("_aws");

assertThat(aws.get("CloudWatchMetrics"))
.asString()
.contains("Namespace=GlobalName");
});
}
}

private void testLogger(Consumer<Consumer<MetricsLogger>> methodToTest) {
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class);
MockedStatic<software.amazon.lambda.powertools.core.internal.SystemWrapper> internalWrapper = mockStatic(
software.amazon.lambda.powertools.core.internal.SystemWrapper.class)) {
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
mocked.when(() -> SystemWrapper.getenv("POWERTOOLS_METRICS_NAMESPACE")).thenReturn("GlobalName");
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID"))
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
.thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");

methodToTest.accept(metricsLogger ->
{
Expand Down
6 changes: 3 additions & 3 deletions ...test/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspectTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ public void metricsWithoutColdStart() {

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

MetricsUtils.defaultDimensions(null);
requestHandler = new PowertoolsMetricsEnabledHandler();
Expand Down Expand Up @@ -135,7 +135,7 @@ public void metricsWithDefaultDimensionSpecified() {

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

requestHandler = new PowertoolsMetricsEnabledDefaultDimensionHandler();

Expand Down Expand Up @@ -179,7 +179,7 @@ public void metricsWithDefaultNoDimensionSpecified() {

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

requestHandler = new PowertoolsMetricsEnabledDefaultNoDimensionHandler();

Expand Down

0 comments on commit 7918d9a

Please sign in to comment.