fix: decouple entry/exit co span classes

Author: tylfin

ddtrace/debugging/_origin/span.py

@@ -199,15 +199,15 @@ def __exit__(self, exc_type, exc_value, traceback):
 class SpanCodeOriginProcessorEntry:
     __uploader__ = LogsIntakeUploaderV1
 
-    _entry_instance: t.Optional["SpanCodeOriginProcessorEntry"] = None
+    _instance: t.Optional["SpanCodeOriginProcessorEntry"] = None
     _handler: t.Optional[t.Callable] = None
 
     @classmethod
     def enable(cls):
-        if cls._entry_instance is not None:
+        if cls._instance is not None:
             return
 
-        cls._entry_instance = cls()
+        cls._instance = cls()
 
         # Register code origin for span with the snapshot uploader
         cls.__uploader__.register(UploaderProduct.CODE_ORIGIN_SPAN)
@@ -218,7 +218,7 @@ def enable(cls):
 
     @classmethod
     def disable(cls):
-        if cls._entry_instance is None:
+        if cls._instance is None:
             return
 
         # Unregister the entrypoint wrapping for entry spans
@@ -227,11 +227,13 @@ def disable(cls):
         cls.__uploader__.unregister(UploaderProduct.CODE_ORIGIN_SPAN)
 
         cls._handler = None
-        cls._entry_instance = None
+        cls._instance = None
 
 
 @dataclass
-class SpanCodeOriginProcessor(SpanCodeOriginProcessorEntry, SpanProcessor):
+class SpanCodeOriginProcessor(SpanProcessor):
+    __uploader__ = LogsIntakeUploaderV1
+
     _instance: t.Optional["SpanCodeOriginProcessor"] = None
 
     def on_span_start(self, span: Span) -> None:
@@ -296,20 +298,19 @@ def enable(cls):
 
         instance = cls._instance = cls()
 
+        # Register code origin for span with the snapshot uploader
+        cls.__uploader__.register(UploaderProduct.CODE_ORIGIN_SPAN)
+
         # Register the processor for exit spans
         instance.register()
 
-        # Enable entry spans. This is idempotent so multiple calls are safe if CO is enabled
-        # after DI via remote configuration.
-        super().enable()
-
     @classmethod
     def disable(cls):
         if cls._instance is None:
             return
 
-        # Disable entry spans
-        super().disable()
+        # Unregister code origin for span with the snapshot uploader
+        cls.__uploader__.unregister(UploaderProduct.CODE_ORIGIN_SPAN)
 
         # Unregister the processor for exit spans
         cls._instance.unregister()

ddtrace/debugging/_products/code_origin/span.py

@@ -17,7 +17,9 @@ def post_preload():
 def start():
     if config.span.enabled:
         from ddtrace.debugging._origin.span import SpanCodeOriginProcessor
+        from ddtrace.debugging._origin.span import SpanCodeOriginProcessorEntry
 
+        SpanCodeOriginProcessorEntry.enable()
         SpanCodeOriginProcessor.enable()
     # If dynamic instrumentation is enabled, and code origin for spans is not explicitly disabled,
     # we'll enable entry spans only.
@@ -34,7 +36,9 @@ 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
 
+        SpanCodeOriginProcessorEntry.disable()
         SpanCodeOriginProcessor.disable()
     elif product_manager.is_enabled(DI_PRODUCT_KEY):
         from ddtrace.debugging._origin.span import SpanCodeOriginProcessorEntry

tests/debugging/origin/test_span.py

@@ -11,6 +11,14 @@
 from tests.utils import TracerTestCase
 
 
+class MockSpanCodeOriginProcessorEntry(SpanCodeOriginProcessorEntry):
+    __uploader__ = MockLogsIntakeUploaderV1
+
+    @classmethod
+    def get_uploader(cls) -> MockLogsIntakeUploaderV1:
+        return t.cast(MockLogsIntakeUploaderV1, cls.__uploader__._instance)
+
+
 class MockSpanCodeOriginProcessor(SpanCodeOriginProcessor):
     __uploader__ = MockLogsIntakeUploaderV1
 
@@ -25,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")
 
@@ -108,14 +118,10 @@ def entry_call():
     def test_span_origin_entry(self):
         # Disable the processor to avoid interference with the test
         MockSpanCodeOriginProcessor.disable()
-        core.reset_listeners(event_id="service_entrypoint.patch")
 
         def entry_call():
             pass
 
-        # Enable entry spans directly
-        SpanCodeOriginProcessorEntry.__uploader__ = MockLogsIntakeUploaderV1
-        SpanCodeOriginProcessorEntry.enable()
         core.dispatch("service_entrypoint.patch", (entry_call,))
 
         with self.tracer.trace("entry"):