diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py index a40e2fcb..260d491b 100644 --- a/datadog_lambda/cold_start.py +++ b/datadog_lambda/cold_start.py @@ -63,6 +63,19 @@ def get_proactive_init_tag(): ) +def is_managed_instances_mode(): + """ + Checks if the Lambda function is running in managed instances mode. + In managed instances mode, we should not create cold start tracing spans + as the gap between the sandbox initialization and the first + invocation might not be be a great experience. + + Returns: + bool: True if running in managed instances mode, False otherwise + """ + return config.aws_lambda_initialization_type == "lambda-managed-instances" + + class ImportNode(object): def __init__(self, module_name, full_file_path, start_time_ns, end_time_ns=None): self.module_name = module_name @@ -145,6 +158,10 @@ def wrapped_find_spec(*args, **kwargs): def initialize_cold_start_tracing(): + # Skip cold start tracing initialization in managed instances mode + if is_managed_instances_mode(): + return + if is_new_sandbox() and config.cold_start_tracing: from sys import meta_path diff --git a/datadog_lambda/config.py b/datadog_lambda/config.py index eda6b582..9960953e 100644 --- a/datadog_lambda/config.py +++ b/datadog_lambda/config.py @@ -98,6 +98,7 @@ def _resolve_env(self, key, default=None, cast=None, depends_on_tracing=False): integration_test = _get_env("DD_INTEGRATION_TEST", "false", as_bool) aws_lambda_function_name = _get_env("AWS_LAMBDA_FUNCTION_NAME") + aws_lambda_initialization_type = _get_env("AWS_LAMBDA_INITIALIZATION_TYPE") @property def function_name(self): diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 5a32e2ff..b2f83f13 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -16,6 +16,7 @@ is_cold_start, is_proactive_init, is_new_sandbox, + is_managed_instances_mode, ColdStartTracer, ) from datadog_lambda.config import config @@ -338,7 +339,13 @@ def _after(self, event, context): create_dd_dummy_metadata_subsegment( self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY ) - should_trace_cold_start = config.cold_start_tracing and is_new_sandbox() + # Skip creating cold start spans in managed instances mode + # In managed instances, the tracer library handles cold start independently + should_trace_cold_start = ( + config.cold_start_tracing + and is_new_sandbox() + and not is_managed_instances_mode() + ) if should_trace_cold_start: trace_ctx = tracer.current_trace_context() diff --git a/tests/test_cold_start.py b/tests/test_cold_start.py index d75b5f43..858db9a5 100644 --- a/tests/test_cold_start.py +++ b/tests/test_cold_start.py @@ -12,6 +12,52 @@ class TestColdStartTracingSetup(unittest.TestCase): + def test_is_managed_instances_mode_when_set(self): + os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] = "lambda-managed-instances" + self.assertTrue(cold_start.is_managed_instances_mode()) + # Clean up + if "AWS_LAMBDA_INITIALIZATION_TYPE" in os.environ: + del os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] + + def test_is_not_managed_instances_mode_when_not_set(self): + if "AWS_LAMBDA_INITIALIZATION_TYPE" in os.environ: + del os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] + self.assertFalse(cold_start.is_managed_instances_mode()) + + def test_is_not_managed_instances_mode_with_different_value(self): + os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] = "on-demand" + self.assertFalse(cold_start.is_managed_instances_mode()) + # Clean up + if "AWS_LAMBDA_INITIALIZATION_TYPE" in os.environ: + del os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] + + def test_initialize_cold_start_tracing_skips_in_managed_instances(self): + # Set managed instances mode + os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] = "lambda-managed-instances" + os.environ["DD_COLD_START_TRACING"] = "true" + cold_start._cold_start = True + cold_start._lambda_container_initialized = False + + # Reset node stacks and wrapped loaders to get clean state + cold_start.reset_node_stacks() + cold_start.already_wrapped_loaders.clear() + + # Count wrapped loaders before + wrapped_loaders_before = len(cold_start.already_wrapped_loaders) + + # Initialize cold start tracing - should skip in managed instances mode + cold_start.initialize_cold_start_tracing() + + # Verify no loaders were wrapped + wrapped_loaders_after = len(cold_start.already_wrapped_loaders) + self.assertEqual(wrapped_loaders_before, wrapped_loaders_after) + + # Clean up + if "AWS_LAMBDA_INITIALIZATION_TYPE" in os.environ: + del os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] + if "DD_COLD_START_TRACING" in os.environ: + del os.environ["DD_COLD_START_TRACING"] + def test_proactive_init(self): cold_start._cold_start = True cold_start._proactive_initialization = False