Cold start EMF capture feature parity

Author: Pankaj Agrawal

powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/PowertoolsMetrics.java

@@ -10,4 +10,5 @@
 public @interface PowertoolsMetrics {
     String namespace() default "";
     String service() default "";
+    boolean captureColdStart() default false;
 }

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

@@ -1,13 +1,20 @@
 package software.amazon.lambda.powertools.metrics.internal;
 
+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;
 import org.aspectj.lang.annotation.Pointcut;
 import software.amazon.cloudwatchlogs.emf.logger.MetricsLogger;
 import software.amazon.cloudwatchlogs.emf.model.DimensionSet;
+import software.amazon.cloudwatchlogs.emf.model.Unit;
 import software.amazon.lambda.powertools.metrics.PowertoolsMetrics;
 
+import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.coldStartDone;
+import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.extractContext;
+import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isColdStart;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isHandlerMethod;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnRequestHandler;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnStreamHandler;
@@ -34,16 +41,47 @@ public Object around(ProceedingJoinPoint pjp,
 
             MetricsLogger logger = logger();
 
-            logger.setNamespace(!"".equals(powertoolsMetrics.namespace()) ? powertoolsMetrics.namespace() : NAMESPACE);
-            logger.setDimensions(DimensionSet.of("service", !"".equals(powertoolsMetrics.service()) ? powertoolsMetrics.service() : serviceName()));
+            logger.setNamespace(namespace(powertoolsMetrics))
+                    .setDimensions(DimensionSet.of("service", service(powertoolsMetrics)));
+
+            MetricsLogger coldStartLogger = null;
+
+            if (powertoolsMetrics.captureColdStart()
+                    && isColdStart()) {
+
+                Optional<Context> contextOptional = extractContext(pjp);
+
+                if (contextOptional.isPresent()) {
+                    Context context = contextOptional.orElseThrow(() -> new IllegalStateException("Context not found"));
+
+                    coldStartLogger = new MetricsLogger();
+
+                    coldStartLogger.putDimensions(DimensionSet.of("function_name", context.getFunctionName()))
+                            .putDimensions(DimensionSet.of("service", service(powertoolsMetrics)))
+                            .setNamespace(namespace(powertoolsMetrics))
+                            .putMetric("ColdStart", 1, Unit.COUNT);
+                }
+            }
 
             Object proceed = pjp.proceed(proceedArgs);
 
-            logger.flush();
+            if (null != coldStartLogger) {
+                coldStartLogger.flush();
+            }
 
+            coldStartDone();
+            logger.flush();
             return proceed;
         }
 
         return pjp.proceed(proceedArgs);
     }
+
+    private String namespace(PowertoolsMetrics powertoolsMetrics) {
+        return !"".equals(powertoolsMetrics.namespace()) ? powertoolsMetrics.namespace() : NAMESPACE;
+    }
+
+    private String service(PowertoolsMetrics powertoolsMetrics) {
+        return !"".equals(powertoolsMetrics.service()) ? powertoolsMetrics.service() : serviceName();
+    }
 }