span registry

Author: ionmincu

src/uipath/core/tracing/manager.py

@@ -1,7 +1,7 @@
 """Tracing manager for handling tracer implementations and function registry."""
 
 import logging
-from typing import Any, Callable, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 from opentelemetry import context, trace
 from opentelemetry.sdk.trace import ReadableSpan
@@ -10,6 +10,84 @@
 logger = logging.getLogger(__name__)
 
 
+class SpanRegistry:
+    """Registry to track all spans and their parent relationships."""
+
+    def __init__(self):
+        self._spans: Dict[int, ReadableSpan] = {}  # span_id -> span
+        self._parent_map: Dict[int, Optional[int]] = {}  # span_id -> parent_id
+
+    def register_span(self, span: ReadableSpan) -> None:
+        """Register a span and its parent relationship."""
+        span_id = span.get_span_context().span_id
+        parent_id = span.parent.span_id if span.parent else None
+
+        self._spans[span_id] = span
+        self._parent_map[span_id] = parent_id
+
+        parent_str = f"{parent_id:016x}" if parent_id is not None else "None"
+        logger.info(f"Registered span: {span.name} (id: {span_id:016x}, parent: {parent_str})")
+
+    def get_span(self, span_id: int) -> Optional[ReadableSpan]:
+        """Get a span by ID."""
+        return self._spans.get(span_id)
+
+    def get_parent_id(self, span_id: int) -> Optional[int]:
+        """Get the parent ID of a span."""
+        return self._parent_map.get(span_id)
+
+    def calculate_depth(self, span_id: int) -> int:
+        """Calculate the depth of a span in the hierarchy."""
+        depth = 0
+        current_id = span_id
+        visited = set()
+
+        while current_id is not None and current_id not in visited:
+            visited.add(current_id)
+            parent_id = self._parent_map.get(current_id)
+            if parent_id is None:
+                break
+            depth += 1
+            current_id = parent_id
+
+        return depth
+
+    def is_ancestor(self, ancestor_id: int, descendant_id: int) -> bool:
+        """Check if ancestor_id is an ancestor of descendant_id."""
+        current_id = descendant_id
+        visited = set()
+
+        while current_id is not None and current_id not in visited:
+            if current_id == ancestor_id:
+                return True
+            visited.add(current_id)
+            current_id = self._parent_map.get(current_id)
+
+        return False
+
+    def get_chain(self, span_id: int) -> List[int]:
+        """Get the complete ancestor chain for a span (bottom to top)."""
+        chain = []
+        current_id = span_id
+        visited = set()
+
+        while current_id is not None and current_id not in visited:
+            chain.append(current_id)
+            visited.add(current_id)
+            current_id = self._parent_map.get(current_id)
+
+        return chain
+
+    def clear(self) -> None:
+        """Clear all registered spans."""
+        self._spans.clear()
+        self._parent_map.clear()
+
+
+# Global span registry instance
+_span_registry = SpanRegistry()
+
+
 class UiPathTracingManager:
     """Static utility class to manage tracing implementations and decorated functions."""
 
@@ -61,43 +139,6 @@ def get_parent_context() -> context.Context:
         )
         return set_span_in_context(bottom_span)
 
-    @staticmethod
-    def _build_ancestor_chain(span: ReadableSpan, spans_by_id: dict) -> List[ReadableSpan]:
-        """Build a complete ancestor chain by traversing parents.
-
-        Args:
-            span: The starting span
-            spans_by_id: Dictionary to populate with discovered spans
-
-        Returns:
-            List of ancestors from bottom (the span) to top (root)
-        """
-        chain = []
-        current = span
-        visited = set()
-
-        while current:
-            span_id = current.get_span_context().span_id
-
-            # Avoid infinite loops
-            if span_id in visited:
-                break
-            visited.add(span_id)
-
-            chain.append(current)
-            spans_by_id[span_id] = current
-
-            # Move to parent if it exists
-            if current.parent:
-                parent_id = current.parent.span_id
-                # Try to get parent from already known spans
-                current = spans_by_id.get(parent_id)
-                # If not found, we can't traverse further
-            else:
-                break
-
-        return chain
-
     @staticmethod
     def _get_bottom_most_span(
         current_span: ReadableSpan, external_span: ReadableSpan
@@ -114,93 +155,68 @@ def _get_bottom_most_span(
         logger.info("=" * 80)
         logger.info("Determining bottom-most span...")
 
-        current_span_id = current_span.get_span_context().span_id
-        external_span_id = external_span.get_span_context().span_id
+        # Register both spans in the registry
+        _span_registry.register_span(current_span)
+        _span_registry.register_span(external_span)
 
-        # Start by adding external ancestors to our known spans
+        # Also register external ancestors
         external_ancestors = UiPathTracingManager.get_ancestor_spans() or []
-        spans_by_id = {}
         for ancestor in external_ancestors:
-            spans_by_id[ancestor.get_span_context().span_id] = ancestor
+            _span_registry.register_span(ancestor)
 
-        # Build complete ancestor chains for both spans
-        logger.info("Building ancestor chain for current_span...")
-        current_chain = UiPathTracingManager._build_ancestor_chain(current_span, spans_by_id)
-
-        logger.info("Building ancestor chain for external_span...")
-        external_chain = UiPathTracingManager._build_ancestor_chain(external_span, spans_by_id)
+        current_span_id = current_span.get_span_context().span_id
+        external_span_id = external_span.get_span_context().span_id
 
-        # Build a parent lookup map
-        parent_map = {}  # span_id -> parent_id
-        for span_id, span in spans_by_id.items():
-            parent_map[span_id] = span.parent.span_id if span.parent else None
+        # Get chains from registry
+        current_chain = _span_registry.get_chain(current_span_id)
+        external_chain = _span_registry.get_chain(external_span_id)
 
-        logger.info(f"current_span chain length (depth): {len(current_chain)}")
-        logger.info(f"external_span chain length (depth): {len(external_chain)}")
+        logger.info(f"current_span: '{current_span.name}' depth: {_span_registry.calculate_depth(current_span_id)}")
+        logger.info(f"external_span: '{external_span.name}' depth: {_span_registry.calculate_depth(external_span_id)}")
 
         # Log the chains
         logger.info("current_span chain (bottom to top):")
-        for i, span in enumerate(current_chain):
-            span_id = span.get_span_context().span_id
-            parent_id = f"{span.parent.span_id:016x}" if span.parent else "None"
-            logger.info(f"  [{i}] {span.name} (id: {span_id:016x}, parent: {parent_id})")
+        for i, span_id in enumerate(current_chain):
+            span = _span_registry.get_span(span_id)
+            if span:
+                parent_id = _span_registry.get_parent_id(span_id)
+                parent_str = f"{parent_id:016x}" if parent_id is not None else "None"
+                logger.info(f"  [{i}] {span.name} (id: {span_id:016x}, parent: {parent_str})")
 
         logger.info("external_span chain (bottom to top):")
-        for i, span in enumerate(external_chain):
-            span_id = span.get_span_context().span_id
-            parent_id = f"{span.parent.span_id:016x}" if span.parent else "None"
-            logger.info(f"  [{i}] {span.name} (id: {span_id:016x}, parent: {parent_id})")
-
-        # Check if one span is an ancestor of the other by walking parent pointers
-        # Check if external_span is an ancestor of current_span
-        current_is_descendant_of_external = False
-        temp_id = current_span_id
-        depth = 0
-        while temp_id is not None and depth < 100:  # Prevent infinite loops
-            if temp_id == external_span_id:
-                current_is_descendant_of_external = True
-                logger.info(f"current_span is a descendant of external_span (found at depth {depth})")
-                break
-            temp_id = parent_map.get(temp_id)
-            depth += 1
-
-        # Check if current_span is an ancestor of external_span
-        external_is_descendant_of_current = False
-        temp_id = external_span_id
-        depth = 0
-        while temp_id is not None and depth < 100:  # Prevent infinite loops
-            if temp_id == current_span_id:
-                external_is_descendant_of_current = True
-                logger.info(f"external_span is a descendant of current_span (found at depth {depth})")
-                break
-            temp_id = parent_map.get(temp_id)
-            depth += 1
-
-        # If they're in the same hierarchy (one is ancestor of the other)
-        if current_is_descendant_of_external:
-            # current_span is deeper (descendant means further down the tree)
+        for i, span_id in enumerate(external_chain):
+            span = _span_registry.get_span(span_id)
+            if span:
+                parent_id = _span_registry.get_parent_id(span_id)
+                parent_str = f"{parent_id:016x}" if parent_id is not None else "None"
+                logger.info(f"  [{i}] {span.name} (id: {span_id:016x}, parent: {parent_str})")
+
+        # Check if one span is an ancestor of the other
+        if _span_registry.is_ancestor(external_span_id, current_span_id):
             logger.info("RESULT: current_span is a descendant of external_span -> returning current_span (deeper)")
             logger.info("=" * 80)
             return current_span
-        elif external_is_descendant_of_current:
-            # external_span is deeper (descendant means further down the tree)
+        elif _span_registry.is_ancestor(current_span_id, external_span_id):
             logger.info("RESULT: external_span is a descendant of current_span -> returning external_span (deeper)")
             logger.info("=" * 80)
             return external_span
 
         # Neither is an ancestor of the other - they're in different branches
-        # Use chain length as tiebreaker
-        if len(current_chain) > len(external_chain):
-            logger.info(f"RESULT: Different branches, current_span is deeper (chain length {len(current_chain)} > {len(external_chain)}) -> returning current_span")
+        # Use depth as tiebreaker
+        current_depth = _span_registry.calculate_depth(current_span_id)
+        external_depth = _span_registry.calculate_depth(external_span_id)
+
+        if current_depth > external_depth:
+            logger.info(f"RESULT: Different branches, current_span is deeper (depth {current_depth} > {external_depth}) -> returning current_span")
             logger.info("=" * 80)
             return current_span
-        elif len(external_chain) > len(current_chain):
-            logger.info(f"RESULT: Different branches, external_span is deeper (chain length {len(external_chain)} > {len(current_chain)}) -> returning external_span")
+        elif external_depth > current_depth:
+            logger.info(f"RESULT: Different branches, external_span is deeper (depth {external_depth} > {current_depth}) -> returning external_span")
             logger.info("=" * 80)
             return external_span
         else:
             # Same depth, different branches - default to external
-            logger.info(f"RESULT: Same depth ({len(current_chain)}), different branches -> defaulting to external_span")
+            logger.info(f"RESULT: Same depth ({current_depth}), different branches -> defaulting to external_span")
             logger.info("=" * 80)
             return external_span