Issue 297: Add AwsRequestId to each EMF log. (#301)

Author: msailes

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

@@ -1,9 +1,7 @@
 package software.amazon.lambda.powertools.metrics.internal;
 
 import java.lang.reflect.Field;
-import java.util.Optional;
 
-import com.amazonaws.services.lambda.runtime.Context;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
@@ -25,7 +23,6 @@
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnStreamHandler;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.serviceName;
 import static software.amazon.lambda.powertools.metrics.MetricsUtils.metricsLogger;
-import static software.amazon.lambda.powertools.metrics.MetricsUtils.withSingleMetric;
 
 @Aspect
 public class LambdaMetricsAspect {
@@ -50,7 +47,10 @@ public Object around(ProceedingJoinPoint pjp,
             logger.setNamespace(namespace(metrics))
                     .putDimensions(DimensionSet.of("Service", service(metrics)));
 
-            coldStartSingleMetricIfApplicable(pjp, metrics);
+            extractContext(pjp).ifPresent((context) -> {
+                coldStartSingleMetricIfApplicable(context.getAwsRequestId(), context.getFunctionName(), metrics);
+                logger.putProperty("AwsRequestId", context.getAwsRequestId());
+            });
 
             try {
                 return pjp.proceed(proceedArgs);
@@ -66,20 +66,19 @@ public Object around(ProceedingJoinPoint pjp,
         return pjp.proceed(proceedArgs);
     }
 
-    private void coldStartSingleMetricIfApplicable(final ProceedingJoinPoint pjp,
+    private void coldStartSingleMetricIfApplicable(final String awsRequestId,
+                                                   final String functionName,
                                                    final Metrics metrics) {
         if (metrics.captureColdStart()
                 && isColdStart()) {
-
-            Optional<Context> contextOptional = extractContext(pjp);
-
-            if (contextOptional.isPresent()) {
-                Context context = contextOptional.orElseThrow(() -> new IllegalStateException("Context not found"));
-
-                withSingleMetric("ColdStart", 1, Unit.COUNT, namespace(metrics), (logger) ->
-                        logger.setDimensions(DimensionSet.of("Service", service(metrics), "FunctionName", context.getFunctionName())));
-            }
+                MetricsLogger metricsLogger = new MetricsLogger();
+                metricsLogger.setNamespace(namespace(metrics));
+                metricsLogger.putMetric("ColdStart", 1, Unit.COUNT);
+                metricsLogger.setDimensions(DimensionSet.of("Service", service(metrics), "FunctionName", functionName));
+                metricsLogger.putProperty("AwsRequestId", awsRequestId);
+                metricsLogger.flush();
         }
+
     }
 
     private void validateBeforeFlushingMetrics(Metrics metrics) {
@@ -103,7 +102,6 @@ private String service(Metrics metrics) {
 
     private void validateMetricsAndRefreshOnFailure(Metrics metrics) {
         try {
-
             validateBeforeFlushingMetrics(metrics);
         } catch (ValidationException e){
             refreshMetricsContext();

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

@@ -82,6 +82,7 @@ public void metricsWithoutColdStart() {
                         assertThat(logAsJson)
                                 .containsEntry("Metric1", 1.0)
                                 .containsEntry("Service", "booking")
+                                .containsEntry("AwsRequestId", "123ABC")
                                 .containsKey("_aws");
                     });
         }
@@ -105,6 +106,7 @@ public void metricsWithColdStart() {
                                 .doesNotContainKey("Metric1")
                                 .containsEntry("ColdStart", 1.0)
                                 .containsEntry("Service", "booking")
+                                .containsEntry("AwsRequestId", "123ABC")
                                 .containsKey("_aws");
 
                         logAsJson = readAsJson(s[1]);
@@ -113,6 +115,7 @@ public void metricsWithColdStart() {
                                 .doesNotContainKey("ColdStart")
                                 .containsEntry("Metric1", 1.0)
                                 .containsEntry("Service", "booking")
+                                .containsEntry("AwsRequestId", "123ABC")
                                 .containsKey("_aws");
                     });
         }
@@ -136,6 +139,7 @@ public void noColdStartMetricsWhenColdStartDone() {
                                 .doesNotContainKey("Metric1")
                                 .containsEntry("ColdStart", 1.0)
                                 .containsEntry("Service", "booking")
+                                .containsEntry("AwsRequestId", "123ABC")
                                 .containsKey("_aws");
 
                         logAsJson = readAsJson(s[1]);
@@ -144,6 +148,7 @@ public void noColdStartMetricsWhenColdStartDone() {
                                 .doesNotContainKey("ColdStart")
                                 .containsEntry("Metric1", 1.0)
                                 .containsEntry("Service", "booking")
+                                .containsEntry("AwsRequestId", "123ABC")
                                 .containsKey("_aws");
 
                         logAsJson = readAsJson(s[2]);
@@ -152,6 +157,7 @@ public void noColdStartMetricsWhenColdStartDone() {
                                 .doesNotContainKey("ColdStart")
                                 .containsEntry("Metric1", 1.0)
                                 .containsEntry("Service", "booking")
+                                .containsEntry("AwsRequestId", "123ABC")
                                 .containsKey("_aws");
                     });
         }
@@ -173,6 +179,7 @@ public void metricsWithStreamHandler() throws IOException {
                         assertThat(logAsJson)
                                 .containsEntry("Metric1", 1.0)
                                 .containsEntry("Service", "booking")
+                                .containsEntry("AwsRequestId", "123ABC")
                                 .containsKey("_aws");
                     });
         }
@@ -249,6 +256,7 @@ public void metricsPublishedEvenHandlerThrowsException() {
                         assertThat(logAsJson)
                                 .containsEntry("CoolMetric", 1.0)
                                 .containsEntry("Service", "booking")
+                                .containsEntry("AwsRequestId", "123ABC")
                                 .containsKey("_aws");
                     });
         }
@@ -259,6 +267,7 @@ private void setupContext() {
         when(context.getInvokedFunctionArn()).thenReturn("testArn");
         when(context.getFunctionVersion()).thenReturn("1");
         when(context.getMemoryLimitInMB()).thenReturn(10);
+        when(context.getAwsRequestId()).thenReturn("123ABC");
     }
 
     private Map<String, Object> readAsJson(String s) {