feat: add support for distributed trace chunks

ddapm_test_agent/trace.py

@@ -232,11 +232,30 @@ def copy_trace(t: Trace) -> Trace:
 
 def root_span(t: Trace) -> Span:
     """Return the root span of the trace."""
+    # Follow approach used in Datadog Agent: https://github.com/DataDog/datadog-agent/blob/927f9ca9acf7983b72a4bfbdd7a69132e1da8501/pkg/trace/traceutil/trace.go#L53
+
+    if len(t) == 0:
+        raise ValueError("empty trace: %s" % t)
+
+    # common case optimization to check for span where parent_id is either not
+    # set or set to 0
     for s in t:
         if "parent_id" not in s or s["parent_id"] is None or s["parent_id"] == 0:
             return s
 
-    raise ValueError("root span not found in trace: %s" % t)
+    # collect root spans as those with parents that are not themselves spans in trace
+    span_ids = set(s["span_id"] for s in t)
+    roots = {
+        s["parent_id"]: s
+        for s in t
+        if "parent_id" in s and s["parent_id"] not in span_ids
+    }
+
+    if len(roots) != 1:
+        raise ValueError("single root span not found in trace: %s" % t)
+
+    # return any root candidate
+    return roots.popitem()[1]
 
 
 def trace_id(t: Trace) -> TraceId:

tests/test_trace.py

@@ -8,6 +8,7 @@
 from ddapm_test_agent.trace import dfs_order
 from ddapm_test_agent.trace import root_span
 
+from .trace_utils import random_id
 from .trace_utils import random_trace
 
 
@@ -20,6 +21,16 @@ def test_random_trace():
         assert bfs_order(t)
 
 
+def test_trace_chunk():
+    trace_id = random_id()
+    parent_id = random_id()
+    t = random_trace(10, trace_id=trace_id, parent_id=parent_id)
+    root = root_span(t)
+    assert root
+    assert root.get("trace_id") == trace_id
+    assert root.get("parent_id") == parent_id
+
+
 @pytest.mark.parametrize(
     "content_type, payload",
     [

tests/trace_utils.py

@@ -1,6 +1,7 @@
 from random import Random
 from typing import Any
 from typing import Dict
+from typing import Optional
 
 from ddapm_test_agent.trace import Span
 from ddapm_test_agent.trace import Trace
@@ -58,7 +59,7 @@ def span(rnd: Random = _random, **kwargs: Any) -> Span:
 
     for k in ["trace_id", "span_id"]:
         if k not in kwargs:
-            kwargs[k] = rnd.randint(0, 2**64)
+            kwargs[k] = random_id(rnd)
 
     # Don't assign a parent id by default
     if "parent_id" not in kwargs:
@@ -112,20 +113,32 @@ def _prufers_trace(n: int, rnd: Random = _random) -> Trace:
     return list(dfs_order(spans))
 
 
-def random_trace(nspans: int, rng: Random = _random) -> Trace:
+def random_trace(
+    nspans: int,
+    rng: Random = _random,
+    trace_id: Optional[int] = None,
+    parent_id: Optional[int] = None,
+) -> Trace:
     # TODO:
     #   represent arbitrary random services (subtrees in spans)
     #   resource names (should only be on service entry)
     #   smarter type (should only be on service entry)
     #   sampling decisions
     #   dd_origin?
     assert nspans > 0
-    trace_id = rng.randint(0, 2**64)
+    if not trace_id:
+        trace_id = random_id(rng)
     t = _prufers_trace(nspans, rng)
     root = root_span(t)
+    if parent_id:
+        root["parent_id"] = parent_id
     for s in t:
         if s is not root:
             del s["type"]
             del s["resource"]
         s["trace_id"] = trace_id
     return t
+
+
+def random_id(rng: Random = _random) -> int:
+    return rng.randint(0, 2**64)