fix(co): ensure entrypoint instrumented (#14790)

## Description

We make sure that entrypoints are instrumented to attach code origin
information on entry span when the feature is enabled remotely.

Author: P403n1x87

.gitlab/benchmarks/bp-runner.microbenchmarks.fail-on-breach.yml

@@ -56,15 +56,15 @@ experiments:
           - name: djangosimple-tracer-minimal
             thresholds:
               - execution_time < 17.50 ms
-              - max_rss_usage < 66.00 MB
+              - max_rss_usage < 67.00 MB
           - name: djangosimple-tracer-native
             thresholds:
               - execution_time < 21.75 ms
               - max_rss_usage < 72.50 MB
           - name: djangosimple-tracer-and-profiler
             thresholds:
               - execution_time < 23.50 ms
-              - max_rss_usage < 67.50 MB
+              - max_rss_usage < 68.00 MB
           - name: djangosimple-tracer-no-caches
             thresholds:
               - execution_time < 19.65 ms
@@ -76,7 +76,7 @@ experiments:
           - name: djangosimple-tracer-dont-create-db-spans
             thresholds:
               - execution_time < 21.50 ms
-              - max_rss_usage < 66.00 MB
+              - max_rss_usage < 67.00 MB
           - name: djangosimple-tracer-no-middleware
             thresholds:
               - execution_time < 21.50 ms

ddtrace/debugging/_debugger.py

@@ -334,7 +334,10 @@ def _dd_debugger_hook(self, probe: Probe) -> None:
                 self._probe_registry.set_emitting(probe)
 
             log.debug("[%s][P: %s] Debugger. Report signal %s", os.getpid(), os.getppid(), signal)
-            self.__uploader__.get_collector().push(signal)
+            if (collector := self.__uploader__.get_collector()) is None:
+                log.error("No collector available to push signal %s", signal)
+                return
+            collector.push(signal)
 
         except Exception:
             log.error("Failed to execute probe hook", exc_info=True)
@@ -474,6 +477,7 @@ def _eject_probes(self, probes_to_eject: List[LineProbe]) -> None:
 
     def _probe_wrapping_hook(self, module: ModuleType) -> None:
         probes = self._probe_registry.get_pending(module.__name__)
+        collector = self.__uploader__.get_collector()
         for probe in probes:
             if not isinstance(probe, FunctionLocationMixin):
                 continue
@@ -500,9 +504,13 @@ def _probe_wrapping_hook(self, module: ModuleType) -> None:
                     function,
                 )
             else:
+                if collector is None:
+                    log.error("No signal collector available")
+                    self._probe_registry.set_error(probe, "NoCollector", "No signal collector available")
+                    continue
                 context = DebuggerWrappingContext(
                     function,
-                    collector=self.__uploader__.get_collector(),
+                    collector=collector,
                     registry=self._probe_registry,
                     tracer=self._tracer,
                     probe_meter=self._probe_meter,

ddtrace/debugging/_exception/replay.py

@@ -283,7 +283,11 @@ def _attach_tb_frame_snapshot_to_span(
                 snapshot.do_line()
 
                 # Collect
-                self.__uploader__.get_collector().push(snapshot)
+                if (collector := self.__uploader__.get_collector()) is None:
+                    log.error("No collector available to push exception replay snapshot")
+                    return False
+
+                collector.push(snapshot)
 
                 # Memoize
                 frame.f_locals[SNAPSHOT_KEY] = snapshot_id = snapshot.uuid

ddtrace/debugging/_origin/span.py

@@ -1,5 +1,4 @@
 from dataclasses import dataclass
-from functools import partial
 from itertools import count
 import sys
 from threading import current_thread
@@ -17,36 +16,30 @@
 from ddtrace.debugging._probe.model import LogLineProbe
 from ddtrace.debugging._probe.model import ProbeEvalTiming
 from ddtrace.debugging._session import Session
-from ddtrace.debugging._signal.collector import SignalCollector
 from ddtrace.debugging._signal.snapshot import Snapshot
 from ddtrace.debugging._uploader import SignalUploader
 from ddtrace.debugging._uploader import UploaderProduct
 from ddtrace.ext import EXIT_SPAN_TYPES
-from ddtrace.internal import core
 from ddtrace.internal.compat import Path
+from ddtrace.internal.forksafe import Lock
+from ddtrace.internal.logger import get_logger
 from ddtrace.internal.packages import is_user_code
 from ddtrace.internal.safety import _isinstance
 from ddtrace.internal.wrapping.context import WrappingContext
 from ddtrace.settings.code_origin import config as co_config
 from ddtrace.trace import Span
 
 
+log = get_logger(__name__)
+
+
 def frame_stack(frame: FrameType) -> t.Iterator[FrameType]:
     _frame: t.Optional[FrameType] = frame
     while _frame is not None:
         yield _frame
         _frame = _frame.f_back
 
 
-def wrap_entrypoint(collector: SignalCollector, f: t.Callable) -> None:
-    if not _isinstance(f, FunctionType):
-        return
-
-    _f = t.cast(FunctionType, f)
-    if not EntrySpanWrappingContext.is_wrapped(_f):
-        EntrySpanWrappingContext(collector, _f).wrap()
-
-
 @dataclass
 class EntrySpanProbe(LogFunctionProbe):
     __span_class__ = "entry"
@@ -118,12 +111,13 @@ class EntrySpanLocation:
 
 
 class EntrySpanWrappingContext(WrappingContext):
+    __enabled__ = False
     __priority__ = 199
 
-    def __init__(self, collector: SignalCollector, f: FunctionType) -> None:
+    def __init__(self, uploader: t.Type[SignalUploader], f: FunctionType) -> None:
         super().__init__(f)
 
-        self.collector = collector
+        self.uploader = uploader
 
         filename = str(Path(f.__code__.co_filename).resolve())
         name = f.__qualname__
@@ -139,26 +133,30 @@ def __init__(self, collector: SignalCollector, f: FunctionType) -> None:
     def __enter__(self):
         super().__enter__()
 
-        root = ddtrace.tracer.current_root_span()
-        span = ddtrace.tracer.current_span()
-        location = self.location
-        if root is None or span is None or root.get_tag("_dd.entry_location.file") is not None:
-            return self
+        if self.__enabled__:
+            root = ddtrace.tracer.current_root_span()
+            span = ddtrace.tracer.current_span()
+            location = self.location
+            if root is None or span is None or root.get_tag("_dd.entry_location.file") is not None:
+                return self
 
-        # Add tags to the local root
-        for s in (root, span):
-            s._set_tag_str("_dd.code_origin.type", "entry")
+            # Add tags to the local root
+            for s in (root, span):
+                s._set_tag_str("_dd.code_origin.type", "entry")
 
-            s._set_tag_str("_dd.code_origin.frames.0.file", location.file)
-            s._set_tag_str("_dd.code_origin.frames.0.line", str(location.line))
-            s._set_tag_str("_dd.code_origin.frames.0.type", location.module)
-            s._set_tag_str("_dd.code_origin.frames.0.method", location.name)
+                s._set_tag_str("_dd.code_origin.frames.0.file", location.file)
+                s._set_tag_str("_dd.code_origin.frames.0.line", str(location.line))
+                s._set_tag_str("_dd.code_origin.frames.0.type", location.module)
+                s._set_tag_str("_dd.code_origin.frames.0.method", location.name)
 
-        self.set("start_time", monotonic_ns())
+            self.set("start_time", monotonic_ns())
 
         return self
 
     def _close_signal(self, retval=None, exc_info=(None, None, None)):
+        if not self.__enabled__:
+            return
+
         root = ddtrace.tracer.current_root_span()
         span = ddtrace.tracer.current_span()
         if root is None or span is None:
@@ -167,6 +165,12 @@ def _close_signal(self, retval=None, exc_info=(None, None, None)):
         # Check if we have any level 2 debugging sessions running for the
         # current trace
         if any(s.level >= 2 for s in Session.from_trace(root.context or span.context)):
+            try:
+                start_time: int = self.get("start_time")
+            except KeyError:
+                # Context was not opened
+                return
+
             # Create a snapshot
             snapshot = Snapshot(
                 probe=self.location.probe,
@@ -182,9 +186,10 @@ def _close_signal(self, retval=None, exc_info=(None, None, None)):
             root._set_tag_str("_dd.code_origin.frames.0.snapshot_id", snapshot.uuid)
             span._set_tag_str("_dd.code_origin.frames.0.snapshot_id", snapshot.uuid)
 
-            snapshot.do_exit(retval, exc_info, monotonic_ns() - self.get("start_time"))
+            snapshot.do_exit(retval, exc_info, monotonic_ns() - start_time)
 
-            self.collector.push(snapshot)
+            if (collector := self.uploader.get_collector()) is not None:
+                collector.push(snapshot)
 
     def __return__(self, retval):
         self._close_signal(retval=retval)
@@ -195,42 +200,68 @@ def __exit__(self, exc_type, exc_value, traceback):
         super().__exit__(exc_type, exc_value, traceback)
 
 
-@dataclass
 class SpanCodeOriginProcessorEntry:
     __uploader__ = SignalUploader
+    __context_wrapper__ = EntrySpanWrappingContext
 
     _instance: t.Optional["SpanCodeOriginProcessorEntry"] = None
-    _handler: t.Optional[t.Callable] = None
+
+    _pending: t.List = []
+    _lock = Lock()
+
+    @classmethod
+    def instrument_view(cls, f):
+        if not _isinstance(f, FunctionType):
+            return
+
+        with cls._lock:
+            if cls._instance is None:
+                # Entry span code origin is not enabled, so we defer the
+                # instrumentation
+                cls._pending.append(f)
+                return
+
+        _f = t.cast(FunctionType, f)
+        if not EntrySpanWrappingContext.is_wrapped(_f):
+            log.debug("Patching entrypoint %r for code origin", f)
+            EntrySpanWrappingContext(cls.__uploader__, _f).wrap()
 
     @classmethod
     def enable(cls):
         if cls._instance is not None:
             return
 
-        cls._instance = cls()
+        with cls._lock:
+            cls._instance = cls()
+
+            # Instrument the pending views
+            while cls._pending:
+                cls.instrument_view(cls._pending.pop())
 
         # Register code origin for span with the snapshot uploader
-        cls.__uploader__.register(UploaderProduct.CODE_ORIGIN_SPAN)
+        cls.__uploader__.register(UploaderProduct.CODE_ORIGIN_SPAN_ENTRY)
 
-        # Register the entrypoint wrapping for entry spans
-        cls._handler = handler = partial(wrap_entrypoint, cls.__uploader__.get_collector())
-        core.on("service_entrypoint.patch", handler)
+        # Enable the context wrapper
+        cls.__context_wrapper__.__enabled__ = True
+
+        log.debug("Code Origin for Spans (entry) enabled")
 
     @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)
+        # Disable the context wrapper
+        cls.__context_wrapper__.__enabled__ = False
+
         # Unregister code origin for span with the snapshot uploader
-        cls.__uploader__.unregister(UploaderProduct.CODE_ORIGIN_SPAN)
+        cls.__uploader__.unregister(UploaderProduct.CODE_ORIGIN_SPAN_ENTRY)
 
-        cls._handler = None
         cls._instance = None
 
+        log.debug("Code Origin for Spans (entry) disabled")
+
 
-@dataclass
 class SpanCodeOriginProcessorExit(SpanProcessor):
     __uploader__ = SignalUploader
 
@@ -283,7 +314,8 @@ def on_span_start(self, span: Span) -> None:
                     snapshot.do_line()
 
                     # Collect
-                    self.__uploader__.get_collector().push(snapshot)
+                    if (collector := self.__uploader__.get_collector()) is not None:
+                        collector.push(snapshot)
 
                     # Correlate the snapshot with the span
                     span._set_tag_str(f"_dd.code_origin.frames.{n}.snapshot_id", snapshot.uuid)
@@ -299,11 +331,13 @@ def enable(cls):
         instance = cls._instance = cls()
 
         # Register code origin for span with the snapshot uploader
-        cls.__uploader__.register(UploaderProduct.CODE_ORIGIN_SPAN)
+        cls.__uploader__.register(UploaderProduct.CODE_ORIGIN_SPAN_EXIT)
 
         # Register the processor for exit spans
         instance.register()
 
+        log.debug("Code Origin for Spans (exit) enabled")
+
     @classmethod
     def disable(cls):
         if cls._instance is None:
@@ -313,6 +347,8 @@ def disable(cls):
         cls._instance.unregister()
 
         # Unregister code origin for span with the snapshot uploader
-        cls.__uploader__.unregister(UploaderProduct.CODE_ORIGIN_SPAN)
+        cls.__uploader__.unregister(UploaderProduct.CODE_ORIGIN_SPAN_EXIT)
 
         cls._instance = None
+
+        log.debug("Code Origin for Spans (exit) disabled")

ddtrace/debugging/_products/code_origin/span.py

@@ -1,10 +1,16 @@
 import enum
+from functools import partial
+import typing as t
 
+import ddtrace.internal.core as core
+from ddtrace.internal.logger import get_logger
 from ddtrace.internal.products import manager as product_manager
 from ddtrace.settings._core import ValueSource
 from ddtrace.settings.code_origin import config
 
 
+log = get_logger(__name__)
+
 CO_ENABLED = "DD_CODE_ORIGIN_FOR_SPANS_ENABLED"
 DI_PRODUCT_KEY = "dynamic-instrumentation"
 
@@ -23,18 +29,26 @@ def _start():
 
 
 def start():
-    if config.span.enabled:
+    # We need to instrument the entrypoints on boot because this is the only
+    # time the tracer will notify us of entrypoints being registered.
+    @partial(core.on, "service_entrypoint.patch")
+    def _(f: t.Callable) -> None:
         from ddtrace.debugging._origin.span import SpanCodeOriginProcessorEntry
+
+        SpanCodeOriginProcessorEntry.instrument_view(f)
+
+    log.debug("Registered entrypoint patching hook for code origin for spans")
+
+    if config.span.enabled:
         from ddtrace.debugging._origin.span import SpanCodeOriginProcessorExit
 
-        SpanCodeOriginProcessorEntry.enable()
         SpanCodeOriginProcessorExit.enable()
+
+        _start()
     # 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
-
-        SpanCodeOriginProcessorEntry.enable()
+        _start()
 
 
 def restart(join=False):

ddtrace/debugging/_uploader.py

@@ -32,7 +32,8 @@ class UploaderProduct(str, Enum):
 
     DEBUGGER = "dynamic_instrumentation"
     EXCEPTION_REPLAY = "exception_replay"
-    CODE_ORIGIN_SPAN = "code_origin.span"
+    CODE_ORIGIN_SPAN_ENTRY = "code_origin.span.entry"
+    CODE_ORIGIN_SPAN_EXIT = "code_origin.span.exit"
 
 
 @dataclass
@@ -220,11 +221,8 @@ def online(self) -> None:
     on_shutdown = online
 
     @classmethod
-    def get_collector(cls) -> SignalCollector:
-        if cls._instance is None:
-            raise RuntimeError("No products registered with the uploader")
-
-        return cls._instance._collector
+    def get_collector(cls) -> Optional[SignalCollector]:
+        return cls._instance._collector if cls._instance is not None else None
 
     @classmethod
     def register(cls, product: UploaderProduct) -> None: