feat(ci_visibility): xdist distributed tracing support (#13509)

Co-authored-by: Vítor De Araújo <vitor.dearaujo@datadoghq.com>

Author: gnufede

ddtrace/contrib/internal/pytest/_plugin_v2.py

@@ -102,6 +102,18 @@
 DISABLED_BY_TEST_MANAGEMENT_REASON = "Flaky test is disabled by Datadog"
 
 
+class XdistHooks:
+    @pytest.hookimpl
+    def pytest_configure_node(self, node):
+        main_session_span = InternalTestSession.get_span()
+        if main_session_span:
+            root_span = main_session_span.span_id
+        else:
+            root_span = 0
+
+        node.workerinput["root_span"] = root_span
+
+
 def _handle_itr_should_skip(item, test_id) -> bool:
     """Checks whether a test should be skipped
 
@@ -252,6 +264,9 @@ def pytest_configure(config: pytest_Config) -> None:
                 from ddtrace.contrib.internal.pytest._pytest_bdd_subplugin import _PytestBddSubPlugin
 
                 config.pluginmanager.register(_PytestBddSubPlugin(), "_datadog-pytest-bdd")
+
+            if config.pluginmanager.hasplugin("xdist"):
+                config.pluginmanager.register(XdistHooks())
         else:
             # If the pytest ddtrace plugin is not enabled, we should disable CI Visibility, as it was enabled during
             # pytest_load_initial_conftests
@@ -299,7 +314,27 @@ def pytest_sessionstart(session: pytest.Session) -> None:
 
         InternalTestSession.set_library_capabilities(library_capabilities)
 
-        InternalTestSession.start()
+        extracted_context = None
+        if hasattr(session.config, "workerinput"):
+            from ddtrace._trace.context import Context
+            from ddtrace.constants import USER_KEEP
+
+            received_root_span = session.config.workerinput.get("root_span", "MISSING_SPAN")
+            try:
+                root_span = int(received_root_span)
+                extracted_context = Context(
+                    trace_id=1,
+                    span_id=root_span,  # This span_id here becomes context.span_id for the parent context
+                    sampling_priority=USER_KEEP,
+                )
+            except ValueError:
+                log.debug(
+                    "pytest_sessionstart: Could not convert root_span %s to int",
+                    received_root_span,
+                )
+
+        InternalTestSession.start(extracted_context)
+
         if InternalTestSession.efd_enabled() and not _pytest_version_supports_efd():
             log.warning("Early Flake Detection disabled: pytest version is not supported")
 
@@ -732,7 +767,10 @@ def _pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
     if ModuleCodeCollector.is_installed():
         ModuleCodeCollector.uninstall()
 
-    InternalTestSession.finish(force_finish_children=True)
+    InternalTestSession.finish(
+        force_finish_children=True,
+        override_status=TestStatus.FAIL if session.exitstatus == pytest.ExitCode.TESTS_FAILED else None,
+    )
 
 
 @pytest.hookimpl(hookwrapper=True, tryfirst=True)

ddtrace/ext/test_visibility/api.py

@@ -23,6 +23,7 @@
 from typing import Optional
 from typing import Type
 
+from ddtrace._trace.context import Context
 from ddtrace.ext.test import Status as _TestStatus
 from ddtrace.ext.test_visibility._item_ids import TestId
 from ddtrace.ext.test_visibility._item_ids import TestModuleId
@@ -179,9 +180,9 @@ def discover(
 
     @staticmethod
     @_catch_and_log_exceptions
-    def start():
+    def start(context: Optional[Context] = None):
         log.debug("Starting session")
-        core.dispatch("test_visibility.session.start")
+        core.dispatch("test_visibility.session.start", (context,))
 
     class FinishArgs(NamedTuple):
         force_finish_children: bool

ddtrace/internal/ci_visibility/api/_base.py

@@ -13,6 +13,7 @@
 from typing import TypeVar
 from typing import Union
 
+from ddtrace._trace.context import Context
 from ddtrace.constants import SPAN_KIND
 from ddtrace.ext import SpanTypes
 from ddtrace.ext import test
@@ -177,9 +178,14 @@ def _add_all_tags_to_span(self) -> None:
             except Exception as e:
                 log.debug("Error setting tag %s: %s", tag, e)
 
-    def _start_span(self) -> None:
+    def _start_span(self, context: Optional[Context] = None) -> None:
         # Test items do not use a parent, and are instead their own trace's root span
-        parent_span = self.get_parent_span() if isinstance(self, TestVisibilityParentItem) else None
+        # except if context is passed (for xdist support)
+        parent_span: Optional[Union[Span, Context]] = None
+        if context is not None:
+            parent_span = context
+        elif isinstance(self, TestVisibilityParentItem):
+            parent_span = self.get_parent_span()
 
         self._span = self._tracer._start_span(
             self._operation_name,
@@ -367,7 +373,7 @@ def _telemetry_record_event_finished(self):
         # Telemetry for events created has specific tags for item types
         raise NotImplementedError("This method must be implemented by the subclass")
 
-    def start(self) -> None:
+    def start(self, context: Optional[Context] = None) -> None:
         log.debug("Test Visibility: starting %s", self)
 
         if self.is_started():
@@ -376,8 +382,9 @@ def start(self) -> None:
                 log.warning(error_msg)
                 raise CIVisibilityDataError(error_msg)
             return
+
         self._telemetry_record_event_created()
-        self._start_span()
+        self._start_span(context)
 
     def is_started(self) -> bool:
         return self._span is not None
@@ -578,7 +585,8 @@ def get_status(self) -> Union[TestStatus, SPECIAL_STATUS]:
 
         if children_status_counts[TestStatus.FAIL.value] > 0:
             return TestStatus.FAIL
-        if children_status_counts[TestStatus.SKIP.value] == len(self._children):
+        len_children = len(self._children)
+        if len_children > 0 and children_status_counts[TestStatus.SKIP.value] == len_children:
             return TestStatus.SKIP
         # We can assume the current item passes if not all children are skipped, and there were no failures
         if children_status_counts[TestStatus.FAIL.value] == 0:

ddtrace/internal/ci_visibility/encoder.py

@@ -1,4 +1,5 @@
 import json
+import os
 import threading
 from typing import TYPE_CHECKING  # noqa:F401
 from uuid import uuid4
@@ -78,8 +79,22 @@ def encode(self):
             self._init_buffer()
             return payload, buffer_size
 
+    def _get_parent_session(self, traces):
+        for trace in traces:
+            for span in trace:
+                if span.get_tag(EVENT_TYPE) == SESSION_TYPE and span.parent_id is not None and span.parent_id != 0:
+                    return span.parent_id
+        return 0
+
     def _build_payload(self, traces):
-        normalized_spans = [self._convert_span(span, trace[0].context.dd_origin) for trace in traces for span in trace]
+        new_parent_session_span_id = self._get_parent_session(traces)
+        is_not_xdist_worker = os.getenv("PYTEST_XDIST_WORKER") is None
+        normalized_spans = [
+            self._convert_span(span, trace[0].context.dd_origin, new_parent_session_span_id)
+            for trace in traces
+            for span in trace
+            if (is_not_xdist_worker or span.get_tag(EVENT_TYPE) != SESSION_TYPE)
+        ]
         if not normalized_spans:
             return None
         record_endpoint_payload_events_count(endpoint=ENDPOINT.TEST_CYCLE, count=len(normalized_spans))
@@ -93,8 +108,8 @@ def _build_payload(self, traces):
     def _pack_payload(payload):
         return msgpack_packb(payload)
 
-    def _convert_span(self, span, dd_origin):
-        # type: (Span, str) -> Dict[str, Any]
+    def _convert_span(self, span, dd_origin, new_parent_session_span_id=0):
+        # type: (Span, str, Optional[int]) -> Dict[str, Any]
         sp = JSONEncoderV2._span_to_dict(span)
         sp = JSONEncoderV2._normalize_span(sp)
         sp["type"] = span.get_tag(EVENT_TYPE) or span.span_type
@@ -103,7 +118,7 @@ def _convert_span(self, span, dd_origin):
         sp["metrics"] = dict(sorted(span._metrics.items()))
         if dd_origin is not None:
             sp["meta"].update({"_dd.origin": dd_origin})
-        sp = CIVisibilityEncoderV01._filter_ids(sp)
+        sp = CIVisibilityEncoderV01._filter_ids(sp, new_parent_session_span_id)
 
         version = CIVisibilityEncoderV01.TEST_SUITE_EVENT_VERSION
         if span.get_tag(EVENT_TYPE) == "test":
@@ -117,7 +132,7 @@ def _convert_span(self, span, dd_origin):
         return {"version": version, "type": event_type, "content": sp}
 
     @staticmethod
-    def _filter_ids(sp):
+    def _filter_ids(sp, new_parent_session_span_id=0):
         """
         Remove trace/span/parent IDs if non-test event, move session/module/suite IDs from meta to outer content layer.
         """
@@ -130,7 +145,7 @@ def _filter_ids(sp):
             sp["parent_id"] = int(sp.get("parent_id") or "1")
             sp["span_id"] = int(sp.get("span_id") or "1")
         if sp["meta"].get(EVENT_TYPE) in [SESSION_TYPE, MODULE_TYPE, SUITE_TYPE, SpanTypes.TEST]:
-            test_session_id = sp["meta"].get(SESSION_ID)
+            test_session_id = new_parent_session_span_id or sp["meta"].get(SESSION_ID)
             if test_session_id:
                 sp[SESSION_ID] = int(test_session_id)
                 del sp["meta"][SESSION_ID]
@@ -221,8 +236,9 @@ def _build_payload(self, traces):
             return None
         return b"\r\n".join(self._build_body(data))
 
-    def _convert_span(self, span, dd_origin):
-        # type: (Span, str) -> Dict[str, Any]
+    def _convert_span(self, span, dd_origin, new_parent_session_span_id=0):
+        # type: (Span, str, Optional[int]) -> Dict[str, Any]
+        # DEV: new_parent_session_span_id is unused here, but it is used in super class
         files: Dict[str, Any] = {}
 
         files_struct_tag_value = span.get_struct_tag(COVERAGE_TAG_NAME)

ddtrace/internal/ci_visibility/recorder.py

@@ -14,6 +14,7 @@
 
 import ddtrace
 from ddtrace import config as ddconfig
+from ddtrace._trace.context import Context
 from ddtrace.contrib import trace_utils
 from ddtrace.ext import ci
 from ddtrace.ext import test
@@ -1102,10 +1103,10 @@ def _on_discover_session(discover_args: TestSession.DiscoverArgs) -> None:
 
 
 @_requires_civisibility_enabled
-def _on_start_session() -> None:
+def _on_start_session(context: Optional[Context] = None) -> None:
     log.debug("Handling start session")
     session = CIVisibility.get_session()
-    session.start()
+    session.start(context)
 
 
 @_requires_civisibility_enabled

releasenotes/notes/feat-ci-vis-xdist-distributed-tracing-support-4f0f59e315c9059e.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    CI Visibility: This introduces preliminary support to link children pytest-xdist tests (and test suites and test modules) to their parent sessions, instead of being sent as independent sessions.