diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index 3e65c3fab96..cb190030432 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/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: diff --git a/aws_lambda_powertools/metrics/provider/cold_start.py b/aws_lambda_powertools/metrics/provider/cold_start.py index c6ef67bd787..4d3aeeefd45 100644 --- a/aws_lambda_powertools/metrics/provider/cold_start.py +++ b/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 diff --git a/aws_lambda_powertools/shared/constants.py b/aws_lambda_powertools/shared/constants.py index 34518fab2e1..a68b59a7c0c 100644 --- a/aws_lambda_powertools/shared/constants.py +++ b/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" diff --git a/aws_lambda_powertools/tracing/tracer.py b/aws_lambda_powertools/tracing/tracer.py index f49fa7bfdbc..503a8e71141 100644 --- a/aws_lambda_powertools/tracing/tracer.py +++ b/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) diff --git a/tests/functional/logger/required_dependencies/test_logger.py b/tests/functional/logger/required_dependencies/test_logger.py index 71dfbc27638..a508209b594 100644 --- a/tests/functional/logger/required_dependencies/test_logger.py +++ b/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") diff --git a/tests/functional/metrics/conftest.py b/tests/functional/metrics/conftest.py index 2de3a0087c2..e9b4edb5935 100644 --- a/tests/functional/metrics/conftest.py +++ b/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) diff --git a/tests/unit/test_tracing.py b/tests/unit/test_tracing.py index 98af062494e..425230380ea 100644 --- a/tests/unit/test_tracing.py +++ b/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)