Reuse cold start variable

Author: Pankaj Agrawal

src/main/java/software/aws/lambda/internal/LambdaHandlerProcessor.java

@@ -9,6 +9,8 @@
 import org.aspectj.lang.ProceedingJoinPoint;
 
 public final class LambdaHandlerProcessor {
+    public static Boolean IS_COLD_START = null;
+
     public static boolean isHandlerMethod(ProceedingJoinPoint pjp) {
         return "handleRequest".equals(pjp.getSignature().getName());
     }

src/main/java/software/aws/lambda/logging/internal/LambdaLoggingAspect.java

@@ -23,13 +23,13 @@
 
 import static java.util.Optional.empty;
 import static java.util.Optional.of;
+import static software.aws.lambda.internal.LambdaHandlerProcessor.IS_COLD_START;
 import static software.aws.lambda.internal.LambdaHandlerProcessor.isHandlerMethod;
 import static software.aws.lambda.internal.LambdaHandlerProcessor.placedOnRequestHandler;
 import static software.aws.lambda.internal.LambdaHandlerProcessor.placedOnStreamHandler;
 
 @Aspect
 public final class LambdaLoggingAspect {
-    static Boolean IS_COLD_START = null;
     private static final ObjectMapper mapper = new ObjectMapper();
 
     @Pointcut("@annotation(powerToolsLogging)")
@@ -47,13 +47,15 @@ public Object around(ProceedingJoinPoint pjp,
                     ThreadContext.put("coldStart", null == IS_COLD_START ? "true" : "false");
                 });
 
-        IS_COLD_START = false;
 
         if (powerToolsLogging.logEvent()) {
             proceedArgs = logEvent(pjp);
         }
 
-        return pjp.proceed(proceedArgs);
+        Object proceed = pjp.proceed(proceedArgs);
+
+        IS_COLD_START = false;
+        return proceed;
     }
 
     private Optional<Context> extractContext(ProceedingJoinPoint pjp) {

src/main/java/software/aws/lambda/tracing/internal/LambdaTracingAspect.java

@@ -38,13 +38,15 @@ public Object around(ProceedingJoinPoint pjp,
             segment.putAnnotation("ColdStart", IS_COLD_START == null);
         }
 
-        IS_COLD_START = false;
 
         try {
             Object methodReturn = pjp.proceed(proceedArgs);
             if (powerToolsTracing.captureResponse()) {
                 segment.putMetadata(namespace(powerToolsTracing), pjp.getSignature().getName() + " response", response(pjp, methodReturn));
             }
+
+            IS_COLD_START = false;
+
             return methodReturn;
         } catch (Exception e) {
             if (powerToolsTracing.captureError()) {

src/test/java/software/aws/lambda/logging/internal/LambdaLoggingAspectTest.java

@@ -15,6 +15,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
+import software.aws.lambda.internal.LambdaHandlerProcessor;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.when;
@@ -32,7 +33,7 @@ class LambdaLoggingAspectTest {
     void setUp() {
         initMocks(this);
         ThreadContext.clearAll();
-        LambdaLoggingAspect.IS_COLD_START = null;
+        LambdaHandlerProcessor.IS_COLD_START = null;
         setupContext();
         requestHandler = new PowerLogToolEnabled();
         requestStreamHandler = new PowerLogToolEnabledForStream();