Skip to content

Commit 10d7ff8

Browse files
author
Pankaj Agrawal
committed
feat: capture xray trace id and request id by default as property for emf logs
1 parent e6f7120 commit 10d7ff8

File tree

5 files changed

+62
-16
lines changed

5 files changed

+62
-16
lines changed

‎powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsUtils.java

+31-5
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package software.amazon.lambda.powertools.metrics;
22

3+
import java.util.Optional;
34
import java.util.function.Consumer;
45

56
import software.amazon.cloudwatchlogs.emf.config.SystemWrapper;
@@ -8,6 +9,11 @@
89
import software.amazon.cloudwatchlogs.emf.model.MetricsLoggerHelper;
910
import software.amazon.cloudwatchlogs.emf.model.Unit;
1011

12+
import static java.util.Optional.ofNullable;
13+
import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.getXrayTraceId;
14+
import static software.amazon.lambda.powertools.metrics.internal.LambdaMetricsAspect.REQUEST_ID_PROPERTY;
15+
import static software.amazon.lambda.powertools.metrics.internal.LambdaMetricsAspect.TRACE_ID_PROPERTY;
16+
1117
/**
1218
* A class used to retrieve the instance of the {@code MetricsLogger} used by
1319
* {@code Metrics}.
@@ -31,21 +37,24 @@ public static MetricsLogger metricsLogger() {
3137

3238
/**
3339
* Add and immediately flush a single metric. It will use the default namespace
34-
* specific either on {@link Metrics} annotation or via POWERTOOLS_METRICS_NAMESPACE env var.
40+
* specified either on {@link Metrics} annotation or via POWERTOOLS_METRICS_NAMESPACE env var.
41+
* It by default captures AwsRequestId as property if used together with {@link Metrics} annotation. It will also
42+
* capture XrayTraceId as property if tracing is enabled.
3543
*
3644
* @param name the name of the metric
3745
* @param value the value of the metric
3846
* @param unit the unit type of the metric
3947
* @param logger the MetricsLogger
4048
*/
41-
public static void withSingleMetricOnDefaultNameSpace(final String name,
42-
final double value,
43-
final Unit unit,
44-
final Consumer<MetricsLogger> logger) {
49+
public static void withSingleMetric(final String name,
50+
final double value,
51+
final Unit unit,
52+
final Consumer<MetricsLogger> logger) {
4553
MetricsLogger metricsLogger = new MetricsLogger();
4654
try {
4755
metricsLogger.setNamespace(defaultNameSpace());
4856
metricsLogger.putMetric(name, value, unit);
57+
captureRequestAndTraceId(metricsLogger);
4958
logger.accept(metricsLogger);
5059
} finally {
5160
metricsLogger.flush();
@@ -54,6 +63,8 @@ public static void withSingleMetricOnDefaultNameSpace(final String name,
5463

5564
/**
5665
* Add and immediately flush a single metric.
66+
* It by default captures AwsRequestId as property if used together with {@link Metrics} annotation. It will also
67+
* capture XrayTraceId as property if tracing is enabled.
5768
*
5869
* @param name the name of the metric
5970
* @param value the value of the metric
@@ -70,15 +81,30 @@ public static void withSingleMetric(final String name,
7081
try {
7182
metricsLogger.setNamespace(namespace);
7283
metricsLogger.putMetric(name, value, unit);
84+
captureRequestAndTraceId(metricsLogger);
7385
logger.accept(metricsLogger);
7486
} finally {
7587
metricsLogger.flush();
7688
}
7789
}
7890

91+
private static void captureRequestAndTraceId(MetricsLogger metricsLogger) {
92+
awsRequestId().
93+
ifPresent(requestId -> metricsLogger.putProperty(REQUEST_ID_PROPERTY, requestId));
94+
95+
getXrayTraceId()
96+
.ifPresent(traceId -> metricsLogger.putProperty(TRACE_ID_PROPERTY, traceId));
97+
}
98+
7999
private static String defaultNameSpace() {
80100
MetricsContext context = MetricsLoggerHelper.metricsContext();
81101
return "aws-embedded-metrics".equals(context.getNamespace()) ?
82102
SystemWrapper.getenv("POWERTOOLS_METRICS_NAMESPACE"): context.getNamespace();
83103
}
104+
105+
private static Optional<String> awsRequestId() {
106+
MetricsContext context = MetricsLoggerHelper.metricsContext();
107+
return ofNullable(context.getProperty(REQUEST_ID_PROPERTY))
108+
.map(Object::toString);
109+
}
84110
}

‎powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspect.java

+8-2
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import software.amazon.cloudwatchlogs.emf.model.DimensionSet;
1111
import software.amazon.cloudwatchlogs.emf.model.MetricsContext;
1212
import software.amazon.cloudwatchlogs.emf.model.Unit;
13+
import software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor;
1314
import software.amazon.lambda.powertools.metrics.Metrics;
1415
import software.amazon.lambda.powertools.metrics.ValidationException;
1516

@@ -27,6 +28,8 @@
2728
@Aspect
2829
public class LambdaMetricsAspect {
2930
private static final String NAMESPACE = System.getenv("POWERTOOLS_METRICS_NAMESPACE");
31+
public static final String TRACE_ID_PROPERTY = "XrayTraceId";
32+
public static final String REQUEST_ID_PROPERTY = "AwsRequestId";
3033

3134
@SuppressWarnings({"EmptyMethod"})
3235
@Pointcut("@annotation(metrics)")
@@ -49,9 +52,12 @@ public Object around(ProceedingJoinPoint pjp,
4952

5053
extractContext(pjp).ifPresent((context) -> {
5154
coldStartSingleMetricIfApplicable(context.getAwsRequestId(), context.getFunctionName(), metrics);
52-
logger.putProperty("AwsRequestId", context.getAwsRequestId());
55+
logger.putProperty(REQUEST_ID_PROPERTY, context.getAwsRequestId());
5356
});
5457

58+
LambdaHandlerProcessor.getXrayTraceId()
59+
.ifPresent(traceId -> logger.putProperty(TRACE_ID_PROPERTY, traceId));
60+
5561
try {
5662
return pjp.proceed(proceedArgs);
5763

@@ -75,7 +81,7 @@ && isColdStart()) {
7581
metricsLogger.setNamespace(namespace(metrics));
7682
metricsLogger.putMetric("ColdStart", 1, Unit.COUNT);
7783
metricsLogger.setDimensions(DimensionSet.of("Service", service(metrics), "FunctionName", functionName));
78-
metricsLogger.putProperty("AwsRequestId", awsRequestId);
84+
metricsLogger.putProperty(REQUEST_ID_PROPERTY, awsRequestId);
7985
metricsLogger.flush();
8086
}
8187

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

+12-5
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import static java.util.Collections.*;
1919
import static org.assertj.core.api.Assertions.assertThat;
2020
import static org.mockito.Mockito.mockStatic;
21+
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getenv;
2122

2223
class MetricsLoggerTest {
2324

@@ -44,8 +45,10 @@ static void beforeAll() {
4445

4546
@Test
4647
void singleMetricsCaptureUtility() {
47-
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class)) {
48+
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class);
49+
MockedStatic<software.amazon.lambda.powertools.core.internal.SystemWrapper> internalWrapper = mockStatic(software.amazon.lambda.powertools.core.internal.SystemWrapper.class)) {
4850
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
51+
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID")).thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
4952

5053
MetricsUtils.withSingleMetric("Metric1", 1, Unit.COUNT, "test",
5154
metricsLogger -> metricsLogger.setDimensions(DimensionSet.of("Dimension1", "Value1")));
@@ -57,18 +60,21 @@ void singleMetricsCaptureUtility() {
5760
assertThat(logAsJson)
5861
.containsEntry("Metric1", 1.0)
5962
.containsEntry("Dimension1", "Value1")
60-
.containsKey("_aws");
63+
.containsKey("_aws")
64+
.containsEntry("XrayTraceId", "1-5759e988-bd862e3fe1be46a994272793");
6165
});
6266
}
6367
}
6468

6569
@Test
6670
void singleMetricsCaptureUtilityWithDefaultNameSpace() {
67-
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class)) {
71+
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class);
72+
MockedStatic<software.amazon.lambda.powertools.core.internal.SystemWrapper> internalWrapper = mockStatic(software.amazon.lambda.powertools.core.internal.SystemWrapper.class)) {
6873
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
6974
mocked.when(() -> SystemWrapper.getenv("POWERTOOLS_METRICS_NAMESPACE")).thenReturn("GlobalName");
75+
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID")).thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
7076

71-
MetricsUtils.withSingleMetricOnDefaultNameSpace("Metric1", 1, Unit.COUNT,
77+
MetricsUtils.withSingleMetric("Metric1", 1, Unit.COUNT,
7278
metricsLogger -> metricsLogger.setDimensions(DimensionSet.of("Dimension1", "Value1")));
7379

7480
assertThat(out.toString())
@@ -78,7 +84,8 @@ void singleMetricsCaptureUtilityWithDefaultNameSpace() {
7884
assertThat(logAsJson)
7985
.containsEntry("Metric1", 1.0)
8086
.containsEntry("Dimension1", "Value1")
81-
.containsKey("_aws");
87+
.containsKey("_aws")
88+
.containsEntry("XrayTraceId", "1-5759e988-bd862e3fe1be46a994272793");
8289

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

‎powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/handlers/PowertoolsMetricsEnabledHandler.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import software.amazon.lambda.powertools.metrics.Metrics;
99

1010
import static software.amazon.lambda.powertools.metrics.MetricsUtils.metricsLogger;
11-
import static software.amazon.lambda.powertools.metrics.MetricsUtils.withSingleMetricOnDefaultNameSpace;
11+
import static software.amazon.lambda.powertools.metrics.MetricsUtils.withSingleMetric;
1212

1313
public class PowertoolsMetricsEnabledHandler implements RequestHandler<Object, Object> {
1414

@@ -19,7 +19,7 @@ public Object handleRequest(Object input, Context context) {
1919
metricsLogger.putMetric("Metric1", 1, Unit.BYTES);
2020

2121

22-
withSingleMetricOnDefaultNameSpace("Metric2", 1, Unit.COUNT,
22+
withSingleMetric("Metric2", 1, Unit.COUNT,
2323
log -> log.setDimensions(DimensionSet.of("Dimension1", "Value1")));
2424

2525
return null;

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

+9-2
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import static org.mockito.Mockito.mockStatic;
3737
import static org.mockito.Mockito.when;
3838
import static org.mockito.MockitoAnnotations.openMocks;
39+
import static software.amazon.lambda.powertools.core.internal.SystemWrapper.getenv;
3940

4041
public class LambdaMetricsAspectTest {
4142
@Mock
@@ -71,8 +72,12 @@ void tearDown() {
7172

7273
@Test
7374
public void metricsWithoutColdStart() {
74-
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class)) {
75+
try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class);
76+
MockedStatic<software.amazon.lambda.powertools.core.internal.SystemWrapper> internalWrapper = mockStatic(software.amazon.lambda.powertools.core.internal.SystemWrapper.class)) {
77+
7578
mocked.when(() -> SystemWrapper.getenv("AWS_EMF_ENVIRONMENT")).thenReturn("Lambda");
79+
internalWrapper.when(() -> getenv("_X_AMZN_TRACE_ID")).thenReturn("Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1\"");
80+
7681
requestHandler.handleRequest("input", context);
7782

7883
assertThat(out.toString().split("\n"))
@@ -83,7 +88,9 @@ public void metricsWithoutColdStart() {
8388
assertThat(logAsJson)
8489
.containsEntry("Metric2", 1.0)
8590
.containsEntry("Dimension1", "Value1")
86-
.containsKey("_aws");
91+
.containsKey("_aws")
92+
.containsEntry("XrayTraceId", "1-5759e988-bd862e3fe1be46a994272793")
93+
.containsEntry("AwsRequestId", "123ABC");
8794

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

0 commit comments

Comments
 (0)