feat(general): make logger, tracer and metrics utilities aware of provisioned concurrency (aws-powertools#6324)

Make utilities aware of Provisioned Concurrency

Author: leandrodamascena

aws_lambda_powertools/logging/logger.py

@@ -64,14 +64,22 @@ def _is_cold_start() -> bool:
     bool
         cold start bool value
     """
-    cold_start = False
-
     global is_cold_start
-    if is_cold_start:
-        cold_start = is_cold_start
+
+    initialization_type = os.getenv(constants.LAMBDA_INITIALIZATION_TYPE)
+
+    # Check for Provisioned Concurrency environment
+    # AWS_LAMBDA_INITIALIZATION_TYPE is set when using Provisioned Concurrency
+    if initialization_type == "provisioned-concurrency":
         is_cold_start = False
+        return False
+
+    if not is_cold_start:
+        return False
 
-    return cold_start
+    # This is a cold start - flip the flag and return True
+    is_cold_start = False
+    return True
 
 
 class Logger:

aws_lambda_powertools/metrics/provider/cold_start.py

@@ -1,7 +1,18 @@
 from __future__ import annotations
 
+import os
+
+from aws_lambda_powertools.shared import constants
+
 is_cold_start = True
 
+initialization_type = os.getenv(constants.LAMBDA_INITIALIZATION_TYPE)
+
+# Check for Provisioned Concurrency environment
+# AWS_LAMBDA_INITIALIZATION_TYPE is set when using Provisioned Concurrency
+if initialization_type == "provisioned-concurrency":
+    is_cold_start = False
+
 
 def reset_cold_start_flag():
     global is_cold_start

aws_lambda_powertools/shared/constants.py

@@ -57,6 +57,7 @@
 SAM_LOCAL_ENV: str = "AWS_SAM_LOCAL"
 CHALICE_LOCAL_ENV: str = "AWS_CHALICE_CLI_MODE"
 LAMBDA_FUNCTION_NAME_ENV: str = "AWS_LAMBDA_FUNCTION_NAME"
+LAMBDA_INITIALIZATION_TYPE: str = "AWS_LAMBDA_INITIALIZATION_TYPE"
 
 # Debug constants
 POWERTOOLS_DEV_ENV: str = "POWERTOOLS_DEV"

aws_lambda_powertools/tracing/tracer.py

@@ -30,6 +30,32 @@
 T = TypeVar("T")
 
 
+def _is_cold_start() -> bool:
+    """Verifies whether is cold start
+
+    Returns
+    -------
+    bool
+        cold start bool value
+    """
+    global is_cold_start
+
+    initialization_type = os.getenv(constants.LAMBDA_INITIALIZATION_TYPE)
+
+    # Check for Provisioned Concurrency environment
+    # AWS_LAMBDA_INITIALIZATION_TYPE is set when using Provisioned Concurrency
+    if initialization_type == "provisioned-concurrency":
+        is_cold_start = False
+        return False
+
+    if not is_cold_start:
+        return False
+
+    # This is a cold start - flip the flag and return True
+    is_cold_start = False
+    return True
+
+
 class Tracer:
     """Tracer using AWS-XRay to provide decorators with known defaults for Lambda functions
 
@@ -340,12 +366,9 @@ def decorate(event, context, **kwargs):
 
                     raise
                 finally:
-                    global is_cold_start
+                    cold_start = _is_cold_start()
                     logger.debug("Annotating cold start")
-                    subsegment.put_annotation(key="ColdStart", value=is_cold_start)
-
-                    if is_cold_start:
-                        is_cold_start = False
+                    subsegment.put_annotation(key="ColdStart", value=cold_start)
 
                     if self.service:
                         subsegment.put_annotation(key="Service", value=self.service)

tests/functional/logger/required_dependencies/test_logger.py

@@ -336,6 +336,33 @@ def handler(event, context):
     assert second_log["cold_start"] is False
 
 
+def test_inject_lambda_cold_start_with_provisioned_concurrency(monkeypatch, lambda_context, stdout, service_name):
+
+    # GIVEN Provisioned Concurrency is enabled via AWS_LAMBDA_INITIALIZATION_TYPE environment variable
+    # AND Logger's cold start flag is explicitly set to True (simulating fresh module import)
+    monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "provisioned-concurrency")
+    from aws_lambda_powertools.logging import logger
+
+    logger.is_cold_start = True
+
+    # GIVEN Logger is initialized
+    logger = Logger(service=service_name, stream=stdout)
+
+    # WHEN a lambda function is decorated with logger, and called twice
+    @logger.inject_lambda_context
+    def handler(event, context):
+        logger.info("Hello")
+
+    handler({}, lambda_context)
+    handler({}, lambda_context)
+
+    # THEN cold_start should be False in both invocations
+    # because Provisioned Concurrency environment variable forces cold_start to always be False
+    first_log, second_log = capture_multiple_logging_statements_output(stdout)
+    assert first_log["cold_start"] is False
+    assert second_log["cold_start"] is False
+
+
 def test_logger_append_duplicated(stdout, service_name):
     # GIVEN Logger is initialized with request_id field
     logger = Logger(service=service_name, stream=stdout, request_id="value")

tests/functional/metrics/conftest.py

@@ -7,7 +7,7 @@
     Metrics,
     MetricUnit,
 )
-from aws_lambda_powertools.metrics.provider.cold_start import reset_cold_start_flag
+from aws_lambda_powertools.metrics.base import reset_cold_start_flag
 
 
 @pytest.fixture(scope="function", autouse=True)

tests/unit/test_tracing.py

@@ -630,6 +630,40 @@ def handler(event, context):
     assert in_subsegment_mock.put_annotation.call_args_list[2] == mocker.call(key="ColdStart", value=False)
 
 
+def test_tracer_lambda_handler_cold_start_with_provisioned_concurrency(
+    monkeypatch,
+    mocker,
+    dummy_response,
+    provider_stub,
+    in_subsegment_mock,
+):
+    # GIVEN Provisioned Concurrency is enabled via AWS_LAMBDA_INITIALIZATION_TYPE environment variable
+    monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "provisioned-concurrency")
+    # GIVEN
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN a Lambda handler is decorated with capture_lambda_handler
+    # AND the handler is invoked twice consecutively
+    @tracer.capture_lambda_handler
+    def handler(event, context):
+        return dummy_response
+
+    # First invocation
+    handler({}, mocker.MagicMock())
+
+    # THEN the ColdStart annotation should be set to False for the first invocation
+    # because Provisioned Concurrency forces cold start to be false regardless of actual state
+    assert in_subsegment_mock.put_annotation.call_args_list[0] == mocker.call(key="ColdStart", value=False)
+
+    # WHEN the same handler is invoked a second time
+    handler({}, mocker.MagicMock())
+
+    # THEN the ColdStart annotation should also be False for the second invocation
+    # confirming that Provisioned Concurrency consistently overrides cold start detection
+    assert in_subsegment_mock.put_annotation.call_args_list[2] == mocker.call(key="ColdStart", value=False)
+
+
 def test_tracer_lambda_handler_add_service_annotation(mocker, dummy_response, provider_stub, in_subsegment_mock):
     # GIVEN
     provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)