fix: omit creating cold start on LMI (#716)

duncanista · web-flow · commit 88bf2f2a5cae · 2026-02-05T15:00:50.000-05:00
* omit creating cold start on LMI

* set env in config instead of using os to do it
diff --git 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
@@ -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
@@ -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
@@ -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
