From 799a227549d1a80306b3b50ea7956d172a4d8e91 Mon Sep 17 00:00:00 2001 From: Debasmita Date: Fri, 5 Jun 2026 22:34:56 +0530 Subject: [PATCH] fix: introduce thread-safe locks in EventBus handler operations --- agentwatch/core/event_bus.py | 148 +++++++++++++++++++---------------- 1 file changed, 80 insertions(+), 68 deletions(-) diff --git a/agentwatch/core/event_bus.py b/agentwatch/core/event_bus.py index da059464..601a7a98 100644 --- a/agentwatch/core/event_bus.py +++ b/agentwatch/core/event_bus.py @@ -9,6 +9,7 @@ import asyncio import inspect import logging +import threading from collections import defaultdict from collections.abc import Callable, Coroutine from datetime import UTC, datetime @@ -114,6 +115,7 @@ def __init__(self) -> None: self._type_index: dict[EventType, set[str]] = defaultdict(set) self._global_handlers: set[str] = set() # subscribed to all events self._lock = asyncio.Lock() + self._thread_lock = threading.Lock() self._event_log: list[AgentEvent] = [] self._max_log_size = 10_000 self._stats: dict[str, int] = defaultdict(int) @@ -138,24 +140,25 @@ def subscribe( def decorator(fn: AnyHandler) -> AnyHandler: _id = handler_id or f"{fn.__module__}.{fn.__qualname__}" - # CodeRabbit: Clean stale registration if ID is reused - if _id in self._handlers: - self.unsubscribe(_id) + with self._thread_lock: + # Clean stale registration if ID is reused + if _id in self._handlers: + self._unsubscribe_locked(_id) + + is_async = inspect.iscoroutinefunction(fn) + reg = HandlerRegistration( + handler_id=_id, + handler=fn, + event_filter=event_filter, + is_async=is_async, + ) + self._handlers[_id] = reg - is_async = inspect.iscoroutinefunction(fn) - reg = HandlerRegistration( - handler_id=_id, - handler=fn, - event_filter=event_filter, - is_async=is_async, - ) - self._handlers[_id] = reg - - if event_types: - for et in event_types: - self._type_index[et].add(_id) - else: - self._global_handlers.add(_id) + if event_types: + for et in event_types: + self._type_index[et].add(_id) + else: + self._global_handlers.add(_id) logger.debug("Registered handler %s for %s", _id, event_types or "ALL") return fn @@ -182,33 +185,29 @@ def subscribe_fn( """ _id = handler_id or f"{fn.__module__}.{fn.__qualname__}.{id(fn)}" - # CodeRabbit: Clean stale registration if ID is reused - if _id in self._handlers: - self.unsubscribe(_id) - - is_async = inspect.iscoroutinefunction(fn) - reg = HandlerRegistration( - handler_id=_id, - handler=fn, - event_filter=event_filter, - is_async=is_async, - ) - self._handlers[_id] = reg - - if event_types: - for et in event_types: - self._type_index[et].add(_id) - else: - self._global_handlers.add(_id) + with self._thread_lock: + # Clean stale registration if ID is reused + if _id in self._handlers: + self._unsubscribe_locked(_id) - return _id + is_async = inspect.iscoroutinefunction(fn) + reg = HandlerRegistration( + handler_id=_id, + handler=fn, + event_filter=event_filter, + is_async=is_async, + ) + self._handlers[_id] = reg - def unsubscribe(self, handler_id: str) -> None: - """Remove a handler and clear it from all type indexes. + if event_types: + for et in event_types: + self._type_index[et].add(_id) + else: + self._global_handlers.add(_id) - Args: - handler_id: ID returned by :meth:`subscribe_fn` or the decorator. - """ + return _id + + def _unsubscribe_locked(self, handler_id: str) -> None: if handler_id not in self._handlers: return self._handlers.pop(handler_id, None) @@ -216,6 +215,15 @@ def unsubscribe(self, handler_id: str) -> None: for ids in self._type_index.values(): ids.discard(handler_id) + def unsubscribe(self, handler_id: str) -> None: + """Remove a handler and clear it from all type indexes. + + Args: + handler_id: ID returned by :meth:`subscribe_fn` or the decorator. + """ + with self._thread_lock: + self._unsubscribe_locked(handler_id) + async def publish(self, event: AgentEvent) -> None: """Publish an event to all matching handlers. @@ -228,28 +236,29 @@ async def publish(self, event: AgentEvent) -> None: event: The normalized event to fan out and log. """ async with self._lock: - # Log to in-memory buffer - self._event_log.append(event) - if len(self._event_log) > self._max_log_size: - self._event_log = self._event_log[-self._max_log_size :] - - self._stats["total_published"] += 1 - self._stats[f"type.{event.event_type.value}"] += 1 - - # Snapshot handler IDs under the lock so subscribe/unsubscribe - # that happens concurrently does not mutate the set mid-iteration. - handler_ids: set[str] = set() - handler_ids.update(self._global_handlers) - handler_ids.update(self._type_index.get(event.event_type, set())) - handlers_to_dispatch = [ - self._handlers[hid] - for hid in handler_ids - if hid in self._handlers - and not ( - self._handlers[hid].event_filter - and not self._handlers[hid].event_filter.matches(event) - ) - ] + with self._thread_lock: + # Log to in-memory buffer + self._event_log.append(event) + if len(self._event_log) > self._max_log_size: + self._event_log = self._event_log[-self._max_log_size :] + + self._stats["total_published"] += 1 + self._stats[f"type.{event.event_type.value}"] += 1 + + # Snapshot handler IDs under the lock so subscribe/unsubscribe + # that happens concurrently does not mutate the set mid-iteration. + handler_ids: set[str] = set() + handler_ids.update(self._global_handlers) + handler_ids.update(self._type_index.get(event.event_type, set())) + handlers_to_dispatch = [ + self._handlers[hid] + for hid in handler_ids + if hid in self._handlers + and not ( + self._handlers[hid].event_filter + and not self._handlers[hid].event_filter.matches(event) + ) + ] # Dispatch outside the lock to avoid holding it during handler I/O tasks = [self._dispatch(reg, event) for reg in handlers_to_dispatch] @@ -340,7 +349,8 @@ def get_recent_events( Returns: Matching events, most recent first. """ - events = list(reversed(self._event_log)) + with self._thread_lock: + events = list(reversed(self._event_log)) if event_type: events = [e for e in events if e.event_type == event_type] if session_id: @@ -353,13 +363,15 @@ def stats(self) -> dict[str, int]: Returns: Copy of internal stats (e.g. ``total_published``, ``type.*``). """ - result = dict(self._stats) - result["active_subscribers"] = self.handler_count() - return result + with self._thread_lock: + result = dict(self._stats) + result["active_subscribers"] = len(self._handlers) + return result def handler_count(self) -> int: """Return the number of registered handlers.""" - return len(self._handlers) + with self._thread_lock: + return len(self._handlers) # Singleton bus instance