fix: linear accumulation of span processors across job runs (#1487)

Author: cristian-groza

packages/uipath-core/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-core"
-version = "0.5.7"
+version = "0.5.8"
 description = "UiPath Core abstractions"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

packages/uipath-core/src/uipath/core/tracing/trace_manager.py

@@ -3,10 +3,13 @@
 from __future__ import annotations
 
 import contextlib
-from typing import Any, Generator, Optional
+import logging
+import threading
+from typing import Any, ClassVar, Generator, Optional
 
+from opentelemetry import context as context_api
 from opentelemetry import trace
-from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor, TracerProvider
 from opentelemetry.sdk.trace.export import SpanExporter
 from opentelemetry.util._decorator import _AgnosticContextManager
 
@@ -17,22 +20,85 @@
 )
 from uipath.core.tracing.types import UiPathTraceSettings
 
+logger = logging.getLogger(__name__)
+
+
+class _DelegatingSpanProcessor(SpanProcessor):
+    """A span processor that delegates to a mutable list of children.
+
+    Registered once on the global TracerProvider. Children can be added
+    and cleared between jobs without touching the provider's internal state.
+
+    WORKAROUND: This exists because OTel's global TracerProvider is set once
+    and has no public API to remove span processors. When a pod runs multiple
+    jobs, each job creates a new UiPathTraceManager that adds processors to
+    the same provider, causing linear accumulation of HTTP calls.
+    The proper fix is for cli_server to create the TracerProvider and
+    trace manager once, and register processors only at startup rather than
+    per-job.
+    """
+
+    _instance: ClassVar[_DelegatingSpanProcessor | None] = None
+    _init_lock: ClassVar[threading.Lock] = threading.Lock()
+
+    def __init__(self) -> None:
+        self._processors: list[SpanProcessor] = []
+
+    @classmethod
+    def get_instance(cls, provider: TracerProvider) -> _DelegatingSpanProcessor:
+        """Get or create the singleton, registering it on the provider once."""
+        if cls._instance is None:
+            with cls._init_lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+                    provider.add_span_processor(cls._instance)
+        return cls._instance
+
+    def add(self, processor: SpanProcessor) -> None:
+        """Add a child processor."""
+        self._processors.append(processor)
+
+    def clear(self) -> list[SpanProcessor]:
+        """Remove and return all child processors."""
+        removed = self._processors.copy()
+        self._processors.clear()
+        return removed
+
+    def on_start(
+        self, span: Span, parent_context: context_api.Context | None = None
+    ) -> None:
+        for p in self._processors:
+            p.on_start(span, parent_context)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        for p in self._processors:
+            p.on_end(span)
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return all(p.force_flush(timeout_millis) for p in self._processors)
+
+    def shutdown(self) -> None:
+        for p in self._processors:
+            p.shutdown()
+
 
 class UiPathTraceManager:
     """Trace manager.
 
-    NOTE: Instantiate trace manager only once.
+    Uses a single delegating processor on the global TracerProvider so that
+    child processors can be added and removed between jobs without accumulation.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize a trace manager."""
         trace.set_tracer_provider(TracerProvider())
-        # If a previous provider set, reuse it.
+        # If a previous provider was already set, reuse it.
         current_provider = trace.get_tracer_provider()
         assert isinstance(current_provider, TracerProvider), (
             "An incompatible Otel TracerProvider was instantiated. Please check runtime configuration."
         )
         self.tracer_provider: TracerProvider = current_provider
+        self._delegating = _DelegatingSpanProcessor.get_instance(current_provider)
         self.tracer_span_processors: list[SpanProcessor] = []
         self.execution_span_exporter = UiPathRuntimeExecutionSpanExporter()
         self.add_span_exporter(self.execution_span_exporter)
@@ -58,7 +124,7 @@ def add_span_exporter(
                 span_exporter, settings
             )
         self.tracer_span_processors.append(span_processor)
-        self.tracer_provider.add_span_processor(span_processor)
+        self._delegating.add(span_processor)
         return self
 
     def add_span_processor(
@@ -67,7 +133,7 @@ def add_span_processor(
     ) -> UiPathTraceManager:
         """Add a span processor to the tracer provider."""
         self.tracer_span_processors.append(span_processor)
-        self.tracer_provider.add_span_processor(span_processor)
+        self._delegating.add(span_processor)
         return self
 
     def get_execution_spans(
@@ -104,5 +170,26 @@ def flush_spans(self) -> None:
         for span_processor in self.tracer_span_processors:
             span_processor.force_flush()
 
+    def shutdown(self) -> None:
+        """Flush, shutdown, and remove all span processors registered by this manager.
+
+        Removes child processors from the delegating processor so they no longer
+        receive span events. Must be called between jobs when the same process
+        handles multiple executions to prevent linear accumulation.
+        """
+        for processor in self.tracer_span_processors:
+            try:
+                processor.force_flush()
+                processor.shutdown()
+            except Exception:
+                logger.warning(
+                    "Failed to shutdown processor %s",
+                    type(processor).__name__,
+                    exc_info=True,
+                )
+        self._delegating.clear()
+        self.tracer_span_processors.clear()
+        self.execution_span_exporter.clear()
+
 
 __all__ = ["UiPathTraceManager"]

packages/uipath-core/tests/tracing/test_trace_manager.py

@@ -1,9 +1,23 @@
-"""Simple test for runtime factory and executor span capture."""
+"""Tests for UiPathTraceManager"""
+
+from unittest.mock import MagicMock
 
 import pytest
 from opentelemetry import trace
+from opentelemetry.sdk.trace import SpanProcessor
+
+from uipath.core.tracing.trace_manager import (
+    UiPathTraceManager,
+    _DelegatingSpanProcessor,
+)
+
 
-from uipath.core.tracing.trace_manager import UiPathTraceManager
+@pytest.fixture(autouse=True)
+def _reset_delegating_singleton():
+    """Reset the singleton between tests so each test starts clean."""
+    _DelegatingSpanProcessor._instance = None
+    yield
+    _DelegatingSpanProcessor._instance = None
 
 
 @pytest.mark.asyncio
@@ -31,3 +45,160 @@ async def test_multiple_factories_same_executor():
 
     assert spans[1].name == "root-span"
     assert spans[1].attributes == {"execution.id": "test"}
+
+
+class TestDelegatingSpanProcessor:
+    """Tests for _DelegatingSpanProcessor."""
+
+    def test_add_and_clear(self) -> None:
+        dp = _DelegatingSpanProcessor()
+        p1 = MagicMock(spec=SpanProcessor)
+        p2 = MagicMock(spec=SpanProcessor)
+
+        dp.add(p1)
+        dp.add(p2)
+        assert len(dp._processors) == 2
+
+        removed = dp.clear()
+        assert removed == [p1, p2]
+        assert dp._processors == []
+
+    def test_delegates_on_start_and_on_end(self) -> None:
+        dp = _DelegatingSpanProcessor()
+        p1 = MagicMock(spec=SpanProcessor)
+        p2 = MagicMock(spec=SpanProcessor)
+        dp.add(p1)
+        dp.add(p2)
+
+        mock_span = MagicMock()
+        dp.on_start(mock_span, parent_context=None)
+        p1.on_start.assert_called_once_with(mock_span, None)
+        p2.on_start.assert_called_once_with(mock_span, None)
+
+        mock_readable_span = MagicMock()
+        dp.on_end(mock_readable_span)
+        p1.on_end.assert_called_once_with(mock_readable_span)
+        p2.on_end.assert_called_once_with(mock_readable_span)
+
+    def test_clear_stops_delegation(self) -> None:
+        dp = _DelegatingSpanProcessor()
+        p1 = MagicMock(spec=SpanProcessor)
+        dp.add(p1)
+        dp.clear()
+
+        mock_span = MagicMock()
+        dp.on_start(mock_span, parent_context=None)
+        dp.on_end(mock_span)
+        p1.on_start.assert_not_called()
+        p1.on_end.assert_not_called()
+
+    def test_get_instance_returns_singleton(self) -> None:
+        provider = MagicMock()
+        inst1 = _DelegatingSpanProcessor.get_instance(provider)
+        inst2 = _DelegatingSpanProcessor.get_instance(provider)
+        assert inst1 is inst2
+        provider.add_span_processor.assert_called_once_with(inst1)
+
+    def test_force_flush_delegates(self) -> None:
+        dp = _DelegatingSpanProcessor()
+        p1 = MagicMock(spec=SpanProcessor)
+        dp.add(p1)
+
+        result = dp.force_flush(timeout_millis=5000)
+        assert result is True
+        p1.force_flush.assert_called_once_with(5000)
+
+    def test_force_flush_returns_false_on_child_failure(self) -> None:
+        dp = _DelegatingSpanProcessor()
+        p1 = MagicMock(spec=SpanProcessor)
+        p2 = MagicMock(spec=SpanProcessor)
+        p1.force_flush.return_value = True
+        p2.force_flush.return_value = False
+        dp.add(p1)
+        dp.add(p2)
+
+        assert dp.force_flush(timeout_millis=5000) is False
+
+    def test_shutdown_delegates_to_children(self) -> None:
+        dp = _DelegatingSpanProcessor()
+        p1 = MagicMock(spec=SpanProcessor)
+        p2 = MagicMock(spec=SpanProcessor)
+        dp.add(p1)
+        dp.add(p2)
+        dp.shutdown()
+        p1.shutdown.assert_called_once()
+        p2.shutdown.assert_called_once()
+        assert len(dp._processors) == 2
+
+
+class TestTraceManagerShutdown:
+    """Tests for the multi-job accumulation fix."""
+
+    def test_shutdown_clears_processors(self) -> None:
+        tm = UiPathTraceManager()
+        p1 = MagicMock(spec=SpanProcessor)
+        tm.add_span_processor(p1)
+
+        tm.shutdown()
+
+        assert tm.tracer_span_processors == []
+        assert tm._delegating._processors == []
+        p1.force_flush.assert_called_once()
+        p1.shutdown.assert_called_once()
+
+    def test_successive_managers_do_not_accumulate(self) -> None:
+        """Simulates multiple jobs on the same pod.
+
+        Each job creates a new UiPathTraceManager and adds a processor.
+        After shutdown, the next job's processors should not stack with
+        the previous ones.
+        """
+        processor_counts: list[int] = []
+
+        for _ in range(5):
+            tm = UiPathTraceManager()
+            mock_processor = MagicMock(spec=SpanProcessor)
+            tm.add_span_processor(mock_processor)
+
+            # Count processors visible to the delegator (minus the batch
+            # processor added by __init__ for execution_span_exporter)
+            processor_counts.append(len(tm._delegating._processors))
+            tm.shutdown()
+
+        # Every job should see the same number of processors (2: the
+        # execution batch processor from __init__ + the one we added),
+        # NOT a linearly growing count.
+        assert all(c == processor_counts[0] for c in processor_counts), (
+            f"Processor counts should be constant across jobs, got: {processor_counts}"
+        )
+
+    def test_shutdown_tolerates_processor_error(self) -> None:
+        tm = UiPathTraceManager()
+        bad_processor = MagicMock(spec=SpanProcessor)
+        bad_processor.force_flush.side_effect = RuntimeError("flush failed")
+        tm.add_span_processor(bad_processor)
+
+        # Should not raise
+        tm.shutdown()
+        assert tm.tracer_span_processors == []
+        assert tm._delegating._processors == []
+
+    def test_spans_not_sent_to_old_processors_after_shutdown(self) -> None:
+        """After shutdown, old processors must not receive new span events."""
+        tm1 = UiPathTraceManager()
+        old_processor = MagicMock(spec=SpanProcessor)
+        tm1.add_span_processor(old_processor)
+        tm1.shutdown()
+
+        # Simulate next job
+        tm2 = UiPathTraceManager()
+        new_processor = MagicMock(spec=SpanProcessor)
+        tm2.add_span_processor(new_processor)
+
+        mock_span = MagicMock()
+        tm2._delegating.on_start(mock_span, parent_context=None)
+
+        new_processor.on_start.assert_called_once()
+        old_processor.on_start.assert_not_called()
+
+        tm2.shutdown()