diff --git a/powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaConstants.java b/powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaConstants.java index 69fc1283a..4c4e8e9db 100644 --- a/powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaConstants.java +++ b/powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaConstants.java @@ -28,4 +28,6 @@ private LambdaConstants() { public static final String ROOT_EQUALS = "Root="; public static final String POWERTOOLS_SERVICE_NAME = "POWERTOOLS_SERVICE_NAME"; public static final String SERVICE_UNDEFINED = "service_undefined"; + public static final String AWS_LAMBDA_INITIALIZATION_TYPE = "AWS_LAMBDA_INITIALIZATION_TYPE"; + public static final String ON_DEMAND_INVOCATION_TYPE = "on-demand"; } diff --git a/powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessor.java b/powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessor.java index 393835d1e..15bff15d6 100644 --- a/powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessor.java +++ b/powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessor.java @@ -93,7 +93,14 @@ protected static void resetServiceName() { } public static boolean isColdStart() { - return isColdStart == null; + if (isColdStart != null) { + return isColdStart; + } + + String initType = System.getenv(LambdaConstants.AWS_LAMBDA_INITIALIZATION_TYPE); + isColdStart = LambdaConstants.ON_DEMAND_INVOCATION_TYPE.equals(initType); + + return isColdStart; } public static void coldStartDone() { diff --git a/powertools-common/src/test/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessorTest.java b/powertools-common/src/test/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessorTest.java index 5c6bdc020..0726a9e77 100644 --- a/powertools-common/src/test/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessorTest.java +++ b/powertools-common/src/test/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessorTest.java @@ -216,6 +216,7 @@ void extractContext_notKnownHandler() { } @Test + @SetEnvironmentVariable(key = LambdaConstants.AWS_LAMBDA_INITIALIZATION_TYPE, value = LambdaConstants.ON_DEMAND_INVOCATION_TYPE) void isColdStart() { boolean isColdStart = LambdaHandlerProcessor.isColdStart();