From ad5cd2e86f276eda65fea7039f729a397441b286 Mon Sep 17 00:00:00 2001 From: Michele Ricciardi Date: Wed, 8 Nov 2023 11:16:34 +0100 Subject: [PATCH 1/3] fix: check if XRAY Trace ID is present in System property --- .../core/internal/LambdaConstants.java | 1 + .../core/internal/LambdaHandlerProcessor.java | 8 ++++- .../core/internal/SystemWrapper.java | 4 +++ .../amazon/lambda/powertools/LoggingE2ET.java | 2 ++ .../internal/LambdaLoggingAspectTest.java | 19 ++++++++++ .../powertools/metrics/MetricsLoggerTest.java | 36 +++++++++++++++++++ 6 files changed, 69 insertions(+), 1 deletion(-) 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..d407d903c 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+, 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..0b6ab01e5 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,6 +246,24 @@ 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"; 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..f3440f359 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; @@ -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( From 1e7feccedfa1bff65ca672d7abae9a857fefab2a Mon Sep 17 00:00:00 2001 From: Michele Ricciardi Date: Wed, 8 Nov 2023 11:19:38 +0100 Subject: [PATCH 2/3] chore: remove erroneous extra char in tests --- .../logging/internal/LambdaLoggingAspectTest.java | 2 +- .../lambda/powertools/metrics/MetricsLoggerTest.java | 8 ++++---- .../metrics/internal/LambdaMetricsAspectTest.java | 6 +++--- 3 files changed, 8 insertions(+), 8 deletions(-) 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 0b6ab01e5..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 @@ -270,7 +270,7 @@ void shouldLogxRayTraceIdEnvVarSet() { 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 f3440f359..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 @@ -67,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")); @@ -97,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"))); @@ -124,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"))); @@ -208,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(); From 9eb1e66e66ac296cd9f11201847fc45fd92076e4 Mon Sep 17 00:00:00 2001 From: Michele Ricciardi Date: Wed, 8 Nov 2023 11:31:22 +0100 Subject: [PATCH 3/3] Update LambdaHandlerProcessor.java --- .../lambda/powertools/core/internal/LambdaHandlerProcessor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 d407d903c..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 @@ -95,7 +95,7 @@ public static boolean isSamLocal() { public static Optional getXrayTraceId() { String X_AMZN_TRACE_ID = getenv(LambdaConstants.X_AMZN_TRACE_ID); - // For the Java Lambda 17+, the Trace ID is set as a System Property + // 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); }