Repository: aws-powertools/powertools-lambda-python

Issue:
Feature request: make cold start aware of Provisioned Concurrency
### Use case

Original issue: https://github.com/aws-powertools/powertools-lambda-typescript/issues/3723

When using Provisioned Concurrency (PC) in AWS Lambda, the service initializes and keeps active a number of execution environments so that a portion of the requests can be served by these already-warm environments.

As a Powertools for AWS customer using PC, I'd like the cold start metric, log key, and tracer annotation to correctly reflect whether a request is a cold start or not - which is not currently the case.

### Solution/User Experience

Currently Logger, Metrics, and Tracer are not aware of Provisioned Concurrency (PC) because the state required to report cold starts in all of them is maintained internally and is dependent on the amount of requests seen by the utility.

Lambda provides an environment variable `AWS_LAMBDA_INITIALIZATION_TYPE` that can have three values in the Python runtime: `provisioned-concurrency`, `snap-start` or `on-demand`.

We could use this value to tweak our cold start detection logic to report a non-cold start (aka warm start) whenever PC is enabled.

This should not require any code changes from customers besides upgrading the packages to the latest version.



### Alternative solutions

```markdown

```

### Acknowledgment

- [x] This feature request meets [Powertools for AWS Lambda (Python) Tenets](https://docs.powertools.aws.dev/lambda/python/latest/#tenets)
- [x] Should this be considered in other Powertools for AWS Lambda languages? i.e. [Java](https://github.com/aws-powertools/powertools-lambda-java/), [TypeScript](https://github.com/aws-powertools/powertools-lambda-typescript/), and [.NET](https://github.com/aws-powertools/powertools-lambda-dotnet/)

Discussion:
@github-actions: ### ⚠️COMMENT VISIBILITY WARNING⚠️
This issue is now closed. Please be mindful that future comments are hard for our team to see.

If you need more assistance, please either tag a [team member](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/MAINTAINERS.md#current-maintainers) or open a new issue that references this one.

If you wish to keep having a conversation with other community members under this issue feel free to do so.@github-actions: This is now released under [3.9.0](https://github.com/aws-powertools/powertools-lambda-python/releases/tag/v3.9.0) version!

Solution patch:
diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py
index 3e65c3fab..cb1900304 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 c6ef67bd7..4d3aeeefd 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 34518fab2..a68b59a7c 100644
--- a/aws_lambda_powertools/shared/constants.py
+++ b/aws_lambda_powertools/shared/constants.py
@@ -57,6 +57,7 @@ LAMBDA_TASK_ROOT_ENV: str = "LAMBDA_TASK_ROOT"
 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 f49fa7bfd..503a8e711 100644
--- a/aws_lambda_powertools/tracing/tracer.py
+++ b/aws_lambda_powertools/tracing/tracer.py
@@ -30,6 +30,32 @@ aws_xray_sdk = LazyLoader(constants.XRAY_SDK_MODULE, globals(), constants.XRAY_S
 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 @@ class Tracer:
 
                     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)


Test patch:
diff --git a/tests/functional/logger/required_dependencies/test_logger.py b/tests/functional/logger/required_dependencies/test_logger.py
index 71dfbc276..a508209b5 100644
--- a/tests/functional/logger/required_dependencies/test_logger.py
+++ b/tests/functional/logger/required_dependencies/test_logger.py
@@ -336,6 +336,33 @@ def test_inject_lambda_cold_start(lambda_context, stdout, service_name):
     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 2de3a0087..e9b4edb59 100644
--- a/tests/functional/metrics/conftest.py
+++ b/tests/functional/metrics/conftest.py
@@ -7,7 +7,7 @@ from aws_lambda_powertools.metrics import (
     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 98af06249..425230380 100644
--- a/tests/unit/test_tracing.py
+++ b/tests/unit/test_tracing.py
@@ -630,6 +630,40 @@ def test_tracer_lambda_handler_cold_start(mocker, dummy_response, provider_stub,
     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)