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"