unify span links (#7327)

Author: cbeauchesne

tests/test_distributed.py

@@ -4,8 +4,7 @@
 
 import json
 from utils import features, interfaces, scenarios, slow, weblog
-from utils.docker_fixtures.spec.trace import span_link_trace_id_equals
-from utils.dd_types import DataDogLibrarySpan, LibraryTraceFormat
+from utils.dd_types import DataDogSpanLink
 
 
 @scenarios.trace_propagation_style_w3c
@@ -56,20 +55,20 @@ def test_span_links_from_conflicting_contexts(self):
         trace = [
             span
             for _, trace, span in interfaces.library.get_spans(self.req, full_trace=True)
-            if _retrieve_span_links(span) is not None
+            if len(span.get_span_links()) != 0
             and trace.trace_id_equals(2)
             and span["parent_id"] == 10  # Only fetch the trace that is related to the header extractions
         ]
 
         assert len(trace) == 1
         span = trace[0]
-        links = _retrieve_span_links(span)
+        links = span.get_span_links()
         assert len(links) == 1
         link1 = links[0]
-        assert span_link_trace_id_equals(link1["trace_id"], 2)
-        assert link1["span_id"] == 987654321
-        assert link1["attributes"] == {"reason": "terminated_context", "context_headers": "tracecontext"}
-        assert link1["trace_id_high"] == 1229782938247303441
+        assert link1.trace_id_low == 2
+        assert link1.span_id == 987654321
+        assert link1.attributes == {"reason": "terminated_context", "context_headers": "tracecontext"}
+        assert link1.trace_id_high == 1229782938247303441
 
     """Datadog and tracecontext headers, trace-id does match, Datadog is primary
     context we want to make sure there's no span link since they match"""
@@ -90,7 +89,7 @@ def test_no_span_links_from_nonconflicting_contexts(self):
         trace = [
             span
             for _, _, span in interfaces.library.get_spans(self.req, full_trace=True)
-            if _retrieve_span_links(span) is not None
+            if len(span.get_span_links()) != 0
             and span["trace_id"] == 1
             and span["parent_id"] == 987654321  # Only fetch the trace that is related to the header extractions
         ]
@@ -117,7 +116,7 @@ def test_no_span_links_from_invalid_trace_id(self):
         trace = [
             span
             for _, _, span in interfaces.library.get_spans(self.req, full_trace=True)
-            if _retrieve_span_links(span) is not None
+            if len(span.get_span_links()) != 0
             and span["trace_id"] == 5
             and span["parent_id"] == 987654324  # Only fetch the trace that is related to the header extractions
         ]
@@ -147,7 +146,7 @@ def test_span_links_flags_from_conflicting_contexts(self):
         spans = [
             span
             for _, _, span in interfaces.library.get_spans(self.req, full_trace=True)
-            if _retrieve_span_links(span) is not None
+            if len(span.get_span_links()) != 0
             and span["trace_id"] == 2
             and span["parent_id"] == 987654321  # Only fetch the trace that is related to the header extractions
         ]
@@ -156,10 +155,10 @@ def test_span_links_flags_from_conflicting_contexts(self):
             raise ValueError(f"Expected 1 span, got {len(spans)}")
 
         span = spans[0]
-        span_links = _retrieve_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 2
         link1 = span_links[0]
-        assert link1["flags"] == 1 | TRACECONTEXT_FLAGS_SET
+        assert link1.data["flags"] == 1 | DataDogSpanLink.TRACECONTEXT_FLAGS_SET
 
 
 @scenarios.default
@@ -183,7 +182,7 @@ def test_span_links_omit_tracestate_from_conflicting_contexts(self):
         spans = [
             span
             for _, _, span in interfaces.library.get_spans(self.req, full_trace=True)
-            if _retrieve_span_links(span) is not None
+            if len(span.get_span_links()) != 0
             and span["trace_id"] == 2
             and span["parent_id"] == 987654321  # Only fetch the trace that is related to the header extractions
         ]
@@ -192,60 +191,10 @@ def test_span_links_omit_tracestate_from_conflicting_contexts(self):
             raise ValueError(f"Expected 1 span, got {len(spans)}")
 
         span = spans[0]
-        links = _retrieve_span_links(span)
+        links = span.get_span_links()
         assert len(links) == 1
         link1 = links[0]
-        assert link1.get("tracestate") is None
-
-
-def _normalize_v1_span_link(link: dict) -> dict:
-    """Ensure v1 span link has trace_id_high when trace_id is 128-bit hex string."""
-    link = dict(link)
-    tid = link.get("trace_id")
-    if isinstance(tid, str) and tid.startswith("0x") and len(tid) > 18 and "trace_id_high" not in link:
-        # 128-bit: high 64 bits (first 16 hex chars after 0x)
-        link["trace_id_high"] = int(tid[2:18], 16)
-    return link
-
-
-def _retrieve_span_links(span: DataDogLibrarySpan):
-    if span.trace.format == LibraryTraceFormat.v10:
-        # v1.0: span_links can be at top level or in attributes
-        links = span.raw_span.get("span_links") or span.raw_span.get("attributes", {}).get("_dd.span_links")
-        if links:
-            return [_normalize_v1_span_link(lnk) for lnk in links]
-        return None
-
-    if span.get("span_links") is not None:
-        return span["span_links"]
-
-    if span["meta"].get("_dd.span_links") is not None:
-        # Convert span_links tags into msgpack v0.4 format
-        json_links = json.loads(span["meta"].get("_dd.span_links"))
-        links = []
-        for json_link in json_links:
-            link = {}
-            link["trace_id"] = int(json_link["trace_id"][-16:], base=16)
-            link["span_id"] = int(json_link["span_id"], base=16)
-            if len(json_link["trace_id"]) > 16:
-                link["trace_id_high"] = int(json_link["trace_id"][:16], base=16)
-            if "attributes" in json_link:
-                link["attributes"] = json_link.get("attributes")
-            if "tracestate" in json_link:
-                link["tracestate"] = json_link.get("tracestate")
-            elif "trace_state" in json_link:
-                link["tracestate"] = json_link.get("trace_state")
-            if "flags" in json_link:
-                link["flags"] = json_link.get("flags") | 1 << 31
-            else:
-                link["flags"] = 0
-            links.append(link)
-        return links
-    return None
-
-
-# The Datadog specific tracecontext flags to mark flags are set
-TRACECONTEXT_FLAGS_SET = 1 << 31
+        assert link1.data.get("tracestate") is None
 
 
 @scenarios.default

tests/test_library_conf.py

@@ -2,8 +2,8 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2021 Datadog, Inc.
 import pytest
-from utils import weblog, interfaces, scenarios, features, logger
-from utils.dd_types import DataDogAgentSpan, AgentTraceFormat
+from utils import weblog, interfaces, scenarios, features
+from utils.dd_types import DataDogAgentSpan
 from utils._context.header_tag_vars import (
     CONFIG_COLON_LEADING,
     CONFIG_COLON_TRAILING,
@@ -426,79 +426,6 @@ def test_trace_header_tags(self):
             assert span_meta[RESPONSE_PREFIX + key.lower()] == response_headers[key]
 
 
-# The Datadog specific tracecontext flags to mark flags are set
-TRACECONTEXT_FLAGS_SET = 1 << 31
-
-
-class SpanLink:
-    def __init__(self, data: dict, trace_id: str, trace_id_low: int, trace_id_high: int):
-        self._data = data
-        self.trace_id = trace_id
-        self.trace_id_low = trace_id_low
-        self.trace_id_high = trace_id_high
-
-        self.attributes: dict[str, str] | None = data.get("attributes")
-        self.trace_state: str | None = data.get("tracestate", data.get("trace_state"))
-        self.flags: int = (data["flags"] | TRACECONTEXT_FLAGS_SET) if "flags" in data else 0
-
-        if "span_id" in self._data:  # span_id is a string on base 16
-            self.span_id = int(data["span_id"], base=16)
-        elif "spanID" in self._data:  # spanID is a string on base 10
-            self.span_id = int(data["spanID"])
-        else:
-            raise ValueError(f"No span id exists in span link: {data}")
-
-    @staticmethod
-    def from_span_links(data: dict) -> "SpanLink":
-        return SpanLink(
-            data, trace_id=data["traceID"], trace_id_high=int(data["traceIDHigh"]), trace_id_low=int(data["traceID"])
-        )
-
-    @staticmethod
-    def from_efficient_trace_payload_format(data: dict) -> "SpanLink":
-        trace_id = data["traceID"]
-
-        return SpanLink(
-            data,
-            trace_id=trace_id,
-            trace_id_high=(int(trace_id, 16) >> 64) & 0xFFFFFFFFFFFFFFFF,
-            trace_id_low=int(trace_id, 16) & 0xFFFFFFFFFFFFFFFF,
-        )
-
-    @staticmethod
-    def from_legacy_format(data: dict) -> "SpanLink":
-        trace_id = data["trace_id"]
-
-        return SpanLink(
-            data,
-            trace_id=trace_id,
-            trace_id_high=int(trace_id[:16], base=16),
-            trace_id_low=int(trace_id[-16:], base=16),
-        )
-
-    def __str__(self) -> str:
-        return str(self._data)
-
-    def __repr__(self) -> str:
-        return repr(self._data)
-
-
-def get_span_links(span: DataDogAgentSpan) -> list[SpanLink]:
-    if span.get("spanLinks") is not None:
-        logger.info("Span links are stored inside span.spanLinks")
-        return [SpanLink.from_span_links(data) for data in span.get("spanLinks")]
-
-    if span.trace.format == AgentTraceFormat.efficient_trace_payload_format and span.get("links") is not None:
-        logger.info("Span links are stored inside span.links and trace format is v1")
-        return [SpanLink.from_efficient_trace_payload_format(data) for data in span.get("links")]
-
-    logger.info("Span links are stored inside span.meta['_dd.span_links'] and trace format is legacy")
-    raw = span.meta.get("_dd.span_links", [])
-    raw_deserilialized = json.loads(raw) if isinstance(raw, (str, bytes, bytearray)) else raw
-
-    return [SpanLink.from_legacy_format(data) for data in raw_deserilialized]
-
-
 @scenarios.default
 @features.context_propagation_extract_behavior
 class Test_ExtractBehavior_Default:
@@ -528,7 +455,7 @@ def test_single_tracecontext(self):
         span = spans[0]
         assert trace_id == 1
         assert span.get("parentID") == "1"
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 0, "Expected span links to be absent"
 
         # Test the next outbound span context
@@ -568,7 +495,7 @@ def test_multiple_tracecontexts(self):
         assert span.get("parentID") == "2"
 
         # Test the extracted span links: One span link per conflicting trace context
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 1, "Expected span links to be present"
 
         link = span_links[0]
@@ -619,7 +546,7 @@ def test_single_tracecontext(self):
         # Test the extracted span links: One span link for the incoming (Datadog trace context).
         # In the case that span links are generated for conflicting trace contexts, those span links
         # are not included in the new trace context
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 1, "Expected span links to be present"
 
         # Assert the Datadog (restarted) span link
@@ -670,7 +597,7 @@ def test_multiple_tracecontexts(self):
         # Test the extracted span links: One span link for the incoming (Datadog trace context).
         # In the case that span links are generated for conflicting trace contexts, those span links
         # are not included in the new trace context
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 1, "Expected span links to be present"
 
         # Assert the Datadog (restarted) span link
@@ -720,7 +647,7 @@ def test_multiple_tracecontexts_with_overrides(self):
         # Test the extracted span links: One span link for the incoming (Datadog trace context).
         # In the case that span links are generated for conflicting trace contexts, those span links
         # are not included in the new trace context
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 1, "Expected span links to be present"
 
         # Assert the Datadog (restarted) span link
@@ -756,7 +683,7 @@ def _get_otel_span(self, spans: list) -> DataDogAgentSpan:
                 return span
         # Fallback: return whichever span has span links
         for span in spans:
-            if get_span_links(span):
+            if span.get_span_links():
                 return span
         return spans[0]
 
@@ -783,7 +710,7 @@ def test_single_tracecontext(self):
         assert span.get("traceID") != "1"
         assert span.get("parentID") is None
 
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 1, "Expected span links to be present"
 
         link = span_links[0]
@@ -822,7 +749,7 @@ def test_multiple_tracecontexts(self):
         assert span.get("traceID") != "8687463697196027922"
         assert span.get("parentID") is None
 
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 1, "Expected span links to be present"
 
         link = span_links[0]
@@ -860,7 +787,7 @@ def test_multiple_tracecontexts_with_overrides(self):
         assert span.get("traceID") != "1"
         assert span.get("parentID") is None
 
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 1, "Expected span links to be present"
 
         link = span_links[0]
@@ -904,7 +831,7 @@ def test_single_tracecontext(self):
         span = spans[0]
         assert span.get("traceID") != "1"
         assert span.get("parentID") is None
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 0, "Expected span links to be absent"
 
         # Test the next outbound span context
@@ -944,7 +871,7 @@ def test_multiple_tracecontexts(self):
             span.get("traceID") != "8687463697196027922"  # Lower 64-bits of traceparent
         )
         assert span.get("parentID") is None
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 0, "Expected span links to be absent"
 
         # Test the next outbound span context
@@ -988,7 +915,7 @@ def test_single_tracecontext(self):
         # Test the extracted span links: One span link for the incoming (Datadog trace context).
         # In the case that span links are generated for conflicting trace contexts, those span links
         # are not included in the new trace context
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 1, "Expected span links to be present"
 
         # Assert the Datadog (restarted) span link
@@ -1039,7 +966,7 @@ def test_multiple_tracecontexts(self):
         # Test the extracted span links: One span link for the incoming (Datadog trace context).
         # In the case that span links are generated for conflicting trace contexts, those span links
         # are not included in the new trace context
-        span_links = get_span_links(span)
+        span_links = span.get_span_links()
         assert len(span_links) == 1, "Expected span links to be present"
 
         # Assert the Datadog (restarted) span link

tests/test_v1_payloads.py

@@ -67,21 +67,20 @@ def setup_span_links_present(self):
     def test_span_links_present(self):
         """V1 spans carrying span links expose them at the top level or in attributes"""
         spans_with_links = [
-            span
-            for _, _, span in interfaces.library.get_spans(self.r, full_trace=True)
-            if span.raw_span.get("span_links") or span.raw_span.get("attributes", {}).get("_dd.span_links")
+            span for _, _, span in interfaces.library.get_spans(self.r, full_trace=True) if span.get_span_links()
         ]
         assert len(spans_with_links) >= 1, "Expected at least one span with span links"
 
         link_carrier = spans_with_links[0]
-        links = link_carrier.raw_span.get("span_links") or []
+        assert link_carrier.trace.format == LibraryTraceFormat.v10
+        links = link_carrier.get_span_links()
 
         assert len(links) >= 1
         link = links[0]
 
-        assert isinstance(link["trace_id"], str)
-        assert link["trace_id"].startswith("0x")
-        assert isinstance(link["span_id"], int)
+        assert isinstance(link.data["trace_id"], str)
+        assert link.data["trace_id"].startswith("0x")
+        assert isinstance(link.data["span_id"], int)
 
 
 @scenarios.apm_tracing_efficient_payload

utils/dd_types/__init__.py

@@ -1,5 +1,6 @@
 from ._datadog_agent_trace import DataDogAgentSpan, DataDogAgentTrace, AgentTraceFormat
 from ._datadog_library_trace import DataDogLibraryTrace, DataDogLibrarySpan, LibraryTraceFormat
+from ._datadog_span_link import DataDogSpanLink
 from ._utils import is_same_boolean
 
 __all__ = [
@@ -8,6 +9,7 @@
     "DataDogAgentTrace",
     "DataDogLibrarySpan",
     "DataDogLibraryTrace",
+    "DataDogSpanLink",
     "LibraryTraceFormat",
     "is_same_boolean",
 ]