diff --git a/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaConstants.java b/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaConstants.java index d0f94260b..e64e334c5 100644 --- a/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaConstants.java +++ b/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaConstants.java @@ -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"; diff --git a/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessor.java b/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessor.java index e9e220e41..91c93c3f3 100644 --- a/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessor.java +++ b/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessor.java @@ -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; @@ -93,7 +94,12 @@ public static boolean isSamLocal() { } public static Optional 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, "")); } diff --git a/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/SystemWrapper.java b/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/SystemWrapper.java index 30f72232f..99046ffec 100644 --- a/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/SystemWrapper.java +++ b/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/SystemWrapper.java @@ -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); + } } diff --git a/powertools-e2e-tests/src/test/java/software/amazon/lambda/powertools/LoggingE2ET.java b/powertools-e2e-tests/src/test/java/software/amazon/lambda/powertools/LoggingE2ET.java index f958970d8..b060879d3 100644 --- a/powertools-e2e-tests/src/test/java/software/amazon/lambda/powertools/LoggingE2ET.java +++ b/powertools-e2e-tests/src/test/java/software/amazon/lambda/powertools/LoggingE2ET.java @@ -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[][] { @@ -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 diff --git a/powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspectTest.java b/powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspectTest.java index b78710586..6952fe755 100644 --- a/powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspectTest.java +++ b/powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspectTest.java @@ -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; @@ -245,13 +246,31 @@ void shouldLogServiceNameWhenEnvVarSet() throws IllegalAccessException { .containsEntry("service", "testService"); } + @Test + void shouldLogxRayTraceIdSystemPropertySet() { + String xRayTraceId = "1-5759e988-bd862e3fe1be46a994272793"; + + try (MockedStatic 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 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); diff --git a/powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/MetricsLoggerTest.java b/powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/MetricsLoggerTest.java index 7f234a4d6..16a068849 100644 --- a/powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/MetricsLoggerTest.java +++ b/powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/MetricsLoggerTest.java @@ -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; @@ -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")); @@ -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"))); @@ -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"))); @@ -165,6 +166,41 @@ void shouldThrowExceptionWhenDefaultDimensionIsNull() { .withMessage("Null dimension set not allowed"); } + @Test + void shouldUseTraceIdFromSystemPropertyIfEnvVarNotPresent() { + try (MockedStatic mocked = mockStatic(SystemWrapper.class); + MockedStatic 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 logAsJson = readAsJson(s); + + assertThat(logAsJson) + .containsEntry("Metric1", 1.0) + .containsEntry("Dimension1", "Value1") + .containsKey("_aws") + .containsEntry("xray_trace_id", "1-5759e988-bd862e3fe1be46a994272793"); + + Map aws = (Map) logAsJson.get("_aws"); + + assertThat(aws.get("CloudWatchMetrics")) + .asString() + .contains("Namespace=GlobalName"); + }); + } + } + private void testLogger(Consumer> methodToTest) { try (MockedStatic mocked = mockStatic(SystemWrapper.class); MockedStatic internalWrapper = mockStatic( @@ -172,7 +208,7 @@ private void testLogger(Consumer> methodToTest) { 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 -> { diff --git a/powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspectTest.java b/powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspectTest.java index eaddfa75d..0a735c75e 100644 --- a/powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspectTest.java +++ b/powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspectTest.java @@ -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(); @@ -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(); @@ -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();