From 85f3abeb8d9c5bab79a92a7d9f6050a466be9c91 Mon Sep 17 00:00:00 2001
From: Rey Abolofia <rey.abolofia@datadoghq.com>
Date: Thu, 15 Feb 2024 14:16:35 -0800
Subject: [PATCH 1/3] Properly parent lazily loaded module imports.

---
 datadog_lambda/cold_start.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py
index 9da02e78..6a76c990 100644
--- a/datadog_lambda/cold_start.py
+++ b/datadog_lambda/cold_start.py
@@ -8,6 +8,7 @@
 _cold_start = True
 _proactive_initialization = False
 _lambda_container_initialized = False
+_tracer = None
 
 
 def set_cold_start(init_timestamp_ns):
@@ -18,6 +19,7 @@ def set_cold_start(init_timestamp_ns):
     global _cold_start
     global _lambda_container_initialized
     global _proactive_initialization
+    global _tracer
     if not _lambda_container_initialized:
         now = time.time_ns()
         if (now - init_timestamp_ns) // 1_000_000_000 > 10:
@@ -29,6 +31,7 @@ def set_cold_start(init_timestamp_ns):
         _cold_start = False
         _proactive_initialization = False
     _lambda_container_initialized = True
+    from ddtrace import tracer as _tracer
 
 
 def is_cold_start():
@@ -62,6 +65,9 @@ def __init__(self, module_name, full_file_path, start_time_ns, end_time_ns=None)
         self.start_time_ns = start_time_ns
         self.end_time_ns = end_time_ns
         self.children = []
+        self.context = None
+        if _lambda_container_initialized:
+            self.context = _tracer.context_provider.active()
 
 
 root_nodes: List[ImportNode] = []
@@ -183,7 +189,8 @@ def trace(self, root_nodes: List[ImportNode] = root_nodes):
         cold_start_span = self.create_cold_start_span(cold_start_span_start_time_ns)
         while root_nodes:
             root_node = root_nodes.pop()
-            self.trace_tree(root_node, cold_start_span)
+            parent = root_node.context or cold_start_span
+            self.trace_tree(root_node, parent)
         self.finish_span(cold_start_span, cold_start_span_end_time_ns)
 
     def trace_tree(self, import_node: ImportNode, parent_span):

From 9ccbed95444f778be685d2d0edeaf1f15fde42a5 Mon Sep 17 00:00:00 2001
From: Rey Abolofia <rey.abolofia@datadoghq.com>
Date: Wed, 13 Mar 2024 11:13:50 -0700
Subject: [PATCH 2/3] Test parenting for lazy loaded package.

---
 tests/test_cold_start.py | 46 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 44 insertions(+), 2 deletions(-)

diff --git a/tests/test_cold_start.py b/tests/test_cold_start.py
index 65193e1d..2ce37e7c 100644
--- a/tests/test_cold_start.py
+++ b/tests/test_cold_start.py
@@ -1,10 +1,13 @@
+import os
 import time
 import unittest
-import datadog_lambda.cold_start as cold_start
+
 from sys import modules, meta_path
-import os
 from unittest.mock import MagicMock
 
+import datadog_lambda.cold_start as cold_start
+import datadog_lambda.wrapper as wrapper
+
 
 class TestColdStartTracingSetup(unittest.TestCase):
     def test_proactive_init(self):
@@ -234,3 +237,42 @@ def test_trace_ignore_libs(self):
         self.cold_start_tracer.trace(nodes)
         self.mock_activate.assert_called_once_with(self.mock_trace_ctx)
         self.assertEqual(self.output_spans, ["node_0", "unittest_cold_start"])
+
+
+def test_lazy_loaded_package_imports(monkeypatch):
+
+    spans = []
+
+    def finish(span):
+        spans.append(span)
+
+    monkeypatch.setattr(wrapper.tracer, "_on_span_finish", finish)
+    monkeypatch.setattr(wrapper, "is_new_sandbox", lambda: True)
+    monkeypatch.setattr("datadog_lambda.wrapper.dd_tracing_enabled", True)
+    monkeypatch.setenv(
+        "DD_COLD_START_TRACE_SKIP_LIB", "ddtrace.contrib.logging,datadog_lambda.wrapper"
+    )
+    monkeypatch.setenv("DD_MIN_COLD_START_DURATION", "0")
+
+    @wrapper.datadog_lambda_wrapper
+    def handler(event, context):
+        import tabnanny
+
+    lambda_context = MagicMock()
+    lambda_context.invoked_function_arn = (
+        "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test:1"
+    )
+
+    handler.cold_start_tracing = True
+    handler({}, lambda_context)
+
+    function_span = import_span = None
+    for span in spans:
+        if span.resource == "tabnanny":
+            import_span = span
+        elif span.name == "aws.lambda":
+            function_span = span
+
+    assert function_span is not None
+    assert import_span is not None
+    assert import_span.parent_id == function_span.span_id

From fb52b6714c300938674840d9973597edd3790adf Mon Sep 17 00:00:00 2001
From: Rey Abolofia <rey.abolofia@datadoghq.com>
Date: Wed, 13 Mar 2024 11:58:48 -0700
Subject: [PATCH 3/3] Clear lists instead of updating the reference.

---
 datadog_lambda/cold_start.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py
index 6a76c990..9dcbec23 100644
--- a/datadog_lambda/cold_start.py
+++ b/datadog_lambda/cold_start.py
@@ -76,10 +76,8 @@ def __init__(self, module_name, full_file_path, start_time_ns, end_time_ns=None)
 
 
 def reset_node_stacks():
-    global root_nodes
-    root_nodes = []
-    global import_stack
-    import_stack = []
+    root_nodes.clear()
+    import_stack.clear()
 
 
 def push_node(module_name, file_path):