feat(di): enable code origins for entry spans when di is enabled (#13593)

Enable code origins for entry spans when dynamic instrumentation is
enabled.

This also breaks up code origin entry and exit spans into separate
classes. This include a new processor for entry spans, updates to
configuration and product management, and additional tests to ensure
functionality and coverage.

Refs: DEBUG-3787

Author: tylfin

ddtrace/debugging/_origin/span.py

@@ -196,12 +196,46 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 
 @dataclass
-class SpanCodeOriginProcessor(SpanProcessor):
+class SpanCodeOriginProcessorEntry:
     __uploader__ = LogsIntakeUploaderV1
 
-    _instance: t.Optional["SpanCodeOriginProcessor"] = None
+    _instance: t.Optional["SpanCodeOriginProcessorEntry"] = None
     _handler: t.Optional[t.Callable] = None
 
+    @classmethod
+    def enable(cls):
+        if cls._instance is not None:
+            return
+
+        cls._instance = cls()
+
+        # Register code origin for span with the snapshot uploader
+        cls.__uploader__.register(UploaderProduct.CODE_ORIGIN_SPAN)
+
+        # Register the entrypoint wrapping for entry spans
+        cls._handler = handler = partial(wrap_entrypoint, cls.__uploader__.get_collector())
+        core.on("service_entrypoint.patch", handler)
+
+    @classmethod
+    def disable(cls):
+        if cls._instance is None:
+            return
+
+        # Unregister the entrypoint wrapping for entry spans
+        core.reset_listeners("service_entrypoint.patch", cls._handler)
+        # Unregister code origin for span with the snapshot uploader
+        cls.__uploader__.unregister(UploaderProduct.CODE_ORIGIN_SPAN)
+
+        cls._handler = None
+        cls._instance = None
+
+
+@dataclass
+class SpanCodeOriginProcessorExit(SpanProcessor):
+    __uploader__ = LogsIntakeUploaderV1
+
+    _instance: t.Optional["SpanCodeOriginProcessorExit"] = None
+
     def on_span_start(self, span: Span) -> None:
         if span.span_type not in EXIT_SPAN_TYPES:
             return
@@ -270,19 +304,11 @@ def enable(cls):
         # Register the processor for exit spans
         instance.register()
 
-        # Register the entrypoint wrapping for entry spans
-        cls._handler = handler = partial(wrap_entrypoint, cls.__uploader__.get_collector())
-        core.on("service_entrypoint.patch", handler)
-
     @classmethod
     def disable(cls):
         if cls._instance is None:
             return
 
-        # Unregister the entrypoint wrapping for entry spans
-        core.reset_listeners("service_entrypoint.patch", cls._handler)
-        cls._handler = None
-
         # Unregister the processor for exit spans
         cls._instance.unregister()
 

ddtrace/debugging/_products/code_origin/span.py

@@ -1,6 +1,11 @@
+from ddtrace.internal.products import manager as product_manager
+from ddtrace.settings._core import ValueSource
 from ddtrace.settings.code_origin import config
 
 
+CO_ENABLED = "DD_CODE_ORIGIN_FOR_SPANS_ENABLED"
+DI_PRODUCT_KEY = "dynamic-instrumentation"
+
 # TODO[gab]: Uncomment this when the feature is ready
 # requires = ["tracer"]
 
@@ -11,9 +16,17 @@ def post_preload():
 
 def start():
     if config.span.enabled:
-        from ddtrace.debugging._origin.span import SpanCodeOriginProcessor
+        from ddtrace.debugging._origin.span import SpanCodeOriginProcessorEntry
+        from ddtrace.debugging._origin.span import SpanCodeOriginProcessorExit
+
+        SpanCodeOriginProcessorEntry.enable()
+        SpanCodeOriginProcessorExit.enable()
+    # If dynamic instrumentation is enabled, and code origin for spans is not explicitly disabled,
+    # we'll enable entry spans only.
+    elif product_manager.is_enabled(DI_PRODUCT_KEY) and config.value_source(CO_ENABLED) == ValueSource.DEFAULT:
+        from ddtrace.debugging._origin.span import SpanCodeOriginProcessorEntry
 
-        SpanCodeOriginProcessor.enable()
+        SpanCodeOriginProcessorEntry.enable()
 
 
 def restart(join=False):
@@ -22,9 +35,15 @@ def restart(join=False):
 
 def stop(join=False):
     if config.span.enabled:
-        from ddtrace.debugging._origin.span import SpanCodeOriginProcessor
+        from ddtrace.debugging._origin.span import SpanCodeOriginProcessorEntry
+        from ddtrace.debugging._origin.span import SpanCodeOriginProcessorExit
+
+        SpanCodeOriginProcessorEntry.disable()
+        SpanCodeOriginProcessorExit.disable()
+    elif product_manager.is_enabled(DI_PRODUCT_KEY):
+        from ddtrace.debugging._origin.span import SpanCodeOriginProcessorEntry
 
-        SpanCodeOriginProcessor.disable()
+        SpanCodeOriginProcessorEntry.disable()
 
 
 def at_exit(join=False):

ddtrace/internal/products.py

@@ -228,5 +228,14 @@ def _() -> None:
         else:
             self._do_products()
 
+    def is_enabled(self, product_name: str, enabled_attribute: str = "enabled") -> bool:
+        if (product := self.__products__.get(product_name)) is None:
+            return False
+
+        if (config := getattr(product, "config", None)) is None:
+            return False
+
+        return getattr(config, enabled_attribute, False)
+
 
 manager = ProductManager()

ddtrace/internal/symbol_db/remoteconfig.py

@@ -14,17 +14,9 @@
 from ddtrace.internal.symbol_db.symbols import SymbolDatabaseUploader
 
 
-log = get_logger(__name__)
-
-
-def is_di_enabled() -> bool:
-    if (di_product := product_manager.__products__.get("dynamic-instrumentation")) is None:
-        return False
+DI_PRODUCT_KEY = "dynamic-instrumentation"
 
-    if (di_config := getattr(di_product, "config", None)) is None:
-        return False
-
-    return di_config.enabled
+log = get_logger(__name__)
 
 
 def _rc_callback(data: t.List[Payload], test_tracer=None):
@@ -57,7 +49,7 @@ def _rc_callback(data: t.List[Payload], test_tracer=None):
             log.debug("[PID %d] SymDB: Symbol DB RCM enablement signal received", os.getpid())
             if not SymbolDatabaseUploader.is_installed():
                 try:
-                    SymbolDatabaseUploader.install(shallow=not is_di_enabled())
+                    SymbolDatabaseUploader.install(shallow=not product_manager.is_enabled(DI_PRODUCT_KEY))
                     log.debug("[PID %d] SymDB: Symbol DB uploader installed", os.getpid())
                 except Exception:
                     log.error("[PID %d] SymDB: Failed to install Symbol DB uploader", os.getpid(), exc_info=True)

releasenotes/notes/entry-spans-di-79eb3b63aab0ecbe.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    dynamic instrumentation: [Code Origins for Spans](https://docs.datadoghq.com/tracing/code_origins/) is now
+    automatically enabled when [Dynamic Instrumentation](https://docs.datadoghq.com/dynamic_instrumentation/) is turned
+    on.

tests/debugging/live/test_live_debugger.py

@@ -2,7 +2,7 @@
 import typing as t
 
 import ddtrace
-from ddtrace.debugging._origin.span import SpanCodeOriginProcessor
+from ddtrace.debugging._origin.span import SpanCodeOriginProcessorExit
 from ddtrace.debugging._probe.model import ProbeEvalTiming
 from ddtrace.internal import core
 from tests.debugging.mocking import MockLogsIntakeUploaderV1
@@ -12,7 +12,7 @@
 from tests.utils import TracerTestCase
 
 
-class MockSpanCodeOriginProcessor(SpanCodeOriginProcessor):
+class MockSpanCodeOriginProcessor(SpanCodeOriginProcessorExit):
     __uploader__ = MockLogsIntakeUploaderV1
 
     @classmethod

tests/debugging/origin/test_span.py

@@ -2,15 +2,24 @@
 import typing as t
 
 import ddtrace
-from ddtrace.debugging._origin.span import SpanCodeOriginProcessor
+from ddtrace.debugging._origin.span import SpanCodeOriginProcessorEntry
+from ddtrace.debugging._origin.span import SpanCodeOriginProcessorExit
 from ddtrace.debugging._session import Session
 from ddtrace.ext import SpanTypes
 from ddtrace.internal import core
 from tests.debugging.mocking import MockLogsIntakeUploaderV1
 from tests.utils import TracerTestCase
 
 
-class MockSpanCodeOriginProcessor(SpanCodeOriginProcessor):
+class MockSpanCodeOriginProcessorEntry(SpanCodeOriginProcessorEntry):
+    __uploader__ = MockLogsIntakeUploaderV1
+
+    @classmethod
+    def get_uploader(cls) -> MockLogsIntakeUploaderV1:
+        return t.cast(MockLogsIntakeUploaderV1, cls.__uploader__._instance)
+
+
+class MockSpanCodeOriginProcessor(SpanCodeOriginProcessorExit):
     __uploader__ = MockLogsIntakeUploaderV1
 
     @classmethod
@@ -24,12 +33,14 @@ def setUp(self):
         self.backup_tracer = ddtrace.tracer
         ddtrace.tracer = self.tracer
 
+        MockSpanCodeOriginProcessorEntry.enable()
         MockSpanCodeOriginProcessor.enable()
 
     def tearDown(self):
         ddtrace.tracer = self.backup_tracer
         super(SpanProbeTestCase, self).tearDown()
 
+        MockSpanCodeOriginProcessorEntry.disable()
         MockSpanCodeOriginProcessor.disable()
         core.reset_listeners(event_id="service_entrypoint.patch")
 
@@ -103,3 +114,40 @@ def entry_call():
         # Check that we have complete data
         snapshot_ids_from_span_tags.add(entry_snapshot_id)
         assert snapshot_ids_from_span_tags == snapshot_ids
+
+    def test_span_origin_entry(self):
+        # Disable the processor to avoid interference with the test
+        MockSpanCodeOriginProcessor.disable()
+
+        def entry_call():
+            pass
+
+        core.dispatch("service_entrypoint.patch", (entry_call,))
+
+        with self.tracer.trace("entry"):
+            entry_call()
+            with self.tracer.trace("middle"):
+                with self.tracer.trace("exit", span_type=SpanTypes.HTTP):
+                    pass
+
+        self.assert_span_count(3)
+        entry, middle, _exit = self.get_spans()
+
+        # Check for the expected tags on the entry span
+        assert entry.get_tag("_dd.code_origin.type") == "entry"
+        assert entry.get_tag("_dd.code_origin.frames.0.file") == str(Path(__file__).resolve())
+        assert entry.get_tag("_dd.code_origin.frames.0.line") == str(entry_call.__code__.co_firstlineno)
+        assert entry.get_tag("_dd.code_origin.frames.0.type") == __name__
+        assert (
+            entry.get_tag("_dd.code_origin.frames.0.method")
+            == "SpanProbeTestCase.test_span_origin_entry.<locals>.entry_call"
+        )
+
+        # Check that we don't have span location tags on the middle span
+        assert middle.get_tag("_dd.code_origin.frames.0.file") is None
+        assert middle.get_tag("_dd.code_origin.frames.0.file") is None
+
+        # Check that we also don't have the span location tags on the exit span
+        assert _exit.get_tag("_dd.code_origin.type") is None
+        assert _exit.get_tag("_dd.code_origin.frames.0.file") is None
+        assert _exit.get_tag("_dd.code_origin.frames.0.line") is None

tests/internal/test_products.py

@@ -113,3 +113,28 @@ def test_product_manager_restart():
         os._exit(0)
 
     os.waitpid(pid, 0)
+
+
+def test_product_manager_is_enabled():
+    class ProductWithConfig:
+        def __init__(self, enabled):
+            self.config = type("Config", (), {"enabled": enabled})()
+
+    # Test when product doesn't exist
+    manager = ProductManagerTest({})
+    assert not manager.is_enabled("nonexistent")
+
+    # Test when product exists but has no config
+    product_no_config = BaseProduct()
+    manager = ProductManagerTest({"no_config": product_no_config})
+    assert not manager.is_enabled("no_config")
+
+    # Test when product exists and is enabled
+    enabled_product = ProductWithConfig(True)
+    manager = ProductManagerTest({"enabled": enabled_product})
+    assert manager.is_enabled("enabled")
+
+    # Test when product exists but is disabled
+    disabled_product = ProductWithConfig(False)
+    manager = ProductManagerTest({"disabled": disabled_product})
+    assert not manager.is_enabled("disabled")