Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Mar 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 13 additions & 5 deletions aws_lambda_powertools/logging/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
11 changes: 11 additions & 0 deletions aws_lambda_powertools/metrics/provider/cold_start.py
Original file line number Diff line number Diff line change
@@ -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

Check warning on line 14 in aws_lambda_powertools/metrics/provider/cold_start.py

View check run for this annotation

Codecov / codecov/patch

aws_lambda_powertools/metrics/provider/cold_start.py#L14

Added line #L14 was not covered by tests


def reset_cold_start_flag():
global is_cold_start
Expand Down
1 change: 1 addition & 0 deletions aws_lambda_powertools/shared/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
33 changes: 28 additions & 5 deletions aws_lambda_powertools/tracing/tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down
27 changes: 27 additions & 0 deletions tests/functional/logger/required_dependencies/test_logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
2 changes: 1 addition & 1 deletion tests/functional/metrics/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
34 changes: 34 additions & 0 deletions tests/unit/test_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading