From 47d4b268da74b3ba0499d963b5982eb1986dc67c Mon Sep 17 00:00:00 2001
From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com>
Date: Tue, 22 Sep 2026 18:03:07 -0400
Subject: [PATCH] fix: use DD_SERVICE for cold start spans instead of hardcoded
 aws.lambda

Cold start spans (aws.lambda.load, aws.lambda.import) were always created
with service="aws.lambda". Use DD_SERVICE when set, falling back to
aws.lambda otherwise, so they land in the same service as the function span.

This matches what the Lambda extension already does by rewriting
aws.lambda spans to DD_SERVICE (datadog-lambda-extension#471), and
follows the same change made for the function span in #635.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 datadog_lambda/cold_start.py |  2 +-
 tests/test_cold_start.py     | 58 +++++++++++++++++++++++++++++++++++-
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py
index 260d491b..408eb476 100644
--- a/datadog_lambda/cold_start.py
+++ b/datadog_lambda/cold_start.py
@@ -248,7 +248,7 @@ def start_span(self, span_type, resource, start_time_ns):
             )  # reactivate required after each finish() call
             self.need_to_reactivate_context = False
         span_kwargs = {
-            "service": "aws.lambda",
+            "service": config.service or "aws.lambda",
             "resource": resource,
             "span_type": span_type,
         }
diff --git a/tests/test_cold_start.py b/tests/test_cold_start.py
index 5fc11e2c..f2f4942f 100644
--- a/tests/test_cold_start.py
+++ b/tests/test_cold_start.py
@@ -3,7 +3,7 @@
 import unittest
 
 from sys import modules, meta_path
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import datadog_lambda.cold_start as cold_start
 import datadog_lambda.wrapper as wrapper
@@ -162,6 +162,7 @@ class TestColdStartTracer(unittest.TestCase):
     def setUp(self) -> None:
         mock_tracer = MagicMock()
         self.output_spans = []
+        self.output_services = []
         self.shared_mock_span = MagicMock()
         self.shared_mock_span.current_spans = []
         self.finish_call_count = 0
@@ -176,6 +177,7 @@ def _finish(finish_time_s):
         def _trace(*args, **kwargs):
             module_name = kwargs["resource"]
             self.shared_mock_span.current_spans.append(module_name)
+            self.output_services.append(kwargs["service"])
             return self.shared_mock_span
 
         mock_tracer.trace = _trace
@@ -286,6 +288,20 @@ def test_trace_ignore_libs(self):
         self.mock_activate.assert_called_once_with(self.mock_trace_ctx)
         self.assertEqual(self.output_spans, ["node_0", "unittest_cold_start"])
 
+    @patch("datadog_lambda.config.Config.service", "my-dd-service")
+    def test_trace_uses_dd_service(self):
+        node_0 = cold_start.ImportNode("node_0", None, self.first_node_start_time_ns)
+        node_0.end_time_ns = self.first_node_start_time_ns + self.test_time_unit
+        self.cold_start_tracer.trace([node_0])
+        self.assertEqual(self.output_services, ["my-dd-service", "my-dd-service"])
+
+    @patch("datadog_lambda.config.Config.service", None)
+    def test_trace_falls_back_to_aws_lambda_service(self):
+        node_0 = cold_start.ImportNode("node_0", None, self.first_node_start_time_ns)
+        node_0.end_time_ns = self.first_node_start_time_ns + self.test_time_unit
+        self.cold_start_tracer.trace([node_0])
+        self.assertEqual(self.output_services, ["aws.lambda", "aws.lambda"])
+
 
 def test_lazy_loaded_package_imports(monkeypatch):
     spans = []
@@ -325,3 +341,43 @@ def handler(event, context):
     assert import_span.trace_id == function_span.trace_id
     assert load_span is not None
     assert load_span.trace_id == function_span.trace_id
+
+
+def test_cold_start_spans_use_dd_service(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.config.Config.trace_enabled", True)
+    monkeypatch.setattr("datadog_lambda.config.Config.service", "my-dd-service")
+    monkeypatch.setenv(
+        "DD_COLD_START_TRACE_SKIP_LIB", "ddtrace.contrib.logging,datadog_lambda.wrapper"
+    )
+    monkeypatch.setenv("DD_MIN_COLD_START_DURATION", "0")
+    # ensure the import below is not served from the module cache
+    monkeypatch.delitem(modules, "colorsys", raising=False)
+
+    @wrapper.datadog_lambda_wrapper
+    def handler(event, context):
+        import colorsys  # noqa: F401
+
+    handler({}, get_mock_context())
+
+    function_span = import_span = load_span = None
+    for span in spans:
+        if span.resource == "colorsys":
+            import_span = span
+        elif span.name == "aws.lambda":
+            function_span = span
+        elif span.name == "aws.lambda.load":
+            load_span = span
+
+    assert function_span is not None
+    assert import_span is not None
+    assert load_span is not None
+    assert function_span.service == "my-dd-service"
+    assert import_span.service == "my-dd-service"
+    assert load_span.service == "my-dd-service"