diff --git a/src/agentex/lib/core/tracing/tracing_processor_manager.py b/src/agentex/lib/core/tracing/tracing_processor_manager.py index 5227e891c..a1b6a3a37 100644 --- a/src/agentex/lib/core/tracing/tracing_processor_manager.py +++ b/src/agentex/lib/core/tracing/tracing_processor_manager.py @@ -4,7 +4,7 @@ import logging import threading from typing import TYPE_CHECKING -from threading import Lock +from threading import RLock from agentex.lib.types.tracing import TracingProcessorConfig from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( @@ -36,7 +36,7 @@ def __init__(self): # Cache for processors self.sync_processors: list[SyncTracingProcessor] = [] self.async_processors: list[AsyncTracingProcessor] = [] - self.lock = Lock() + self.lock = RLock() self._agentex_registered = False def _ensure_agentex_registered(self): diff --git a/tests/lib/core/tracing/test_tracing_processor_manager.py b/tests/lib/core/tracing/test_tracing_processor_manager.py new file mode 100644 index 000000000..d17ac6581 --- /dev/null +++ b/tests/lib/core/tracing/test_tracing_processor_manager.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import threading +from typing import Any, cast +from dataclasses import dataclass + +from agentex.lib.core.tracing.tracing_processor_manager import TracingProcessorManager + + +@dataclass +class _FakeConfig: + type: str = "fake" + + +class _FakeProcessor: + def __init__(self, config: Any) -> None: + self.config = config + + +def _manager() -> TracingProcessorManager: + manager = TracingProcessorManager() + manager.sync_config_registry["fake"] = cast(Any, _FakeProcessor) + manager.async_config_registry["fake"] = cast(Any, _FakeProcessor) + return manager + + +REGISTRATION_DEADLINE_SECONDS = 5.0 + + +def _finishes(target: Any, timeout: float = REGISTRATION_DEADLINE_SECONDS) -> bool: + errors: list[BaseException] = [] + + def _run() -> None: + try: + target() + except BaseException as exc: + errors.append(exc) + + thread = threading.Thread(target=_run, daemon=True) + thread.start() + thread.join(timeout) + if errors: + raise errors[0] + return not thread.is_alive() + + +def test_set_processor_configs_registers_every_config_without_deadlocking() -> None: + manager = _manager() + configs = [_FakeConfig(), _FakeConfig()] + + finished = _finishes(lambda: manager.set_processor_configs(cast(Any, configs))) + + assert finished, "set_processor_configs blocked on the manager's own lock" + assert [cast(Any, p).config for p in manager.get_sync_processors()] == configs + assert [cast(Any, p).config for p in manager.get_async_processors()] == configs + + +def test_add_processor_config_still_registers_one_pair() -> None: + manager = _manager() + + manager.add_processor_config(cast(Any, _FakeConfig())) + + assert len(manager.get_sync_processors()) == 1 + assert len(manager.get_async_processors()) == 1