diff --git a/pyproject.toml b/pyproject.toml index a10b02e69b..128098d257 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,6 @@ dependencies = [ "psutil>=6.0.0", "pydantic-settings>=2.12.0", "pydantic>=2.11.0", - "pyee>=9.0.0", "tldextract>=5.3.0", "typing-extensions>=4.10.0", "yarl>=1.18.0", diff --git a/src/crawlee/_utils/wait.py b/src/crawlee/_utils/wait.py index 1c889908e7..8c18c2e494 100644 --- a/src/crawlee/_utils/wait.py +++ b/src/crawlee/_utils/wait.py @@ -46,7 +46,7 @@ async def wait_for( raise RuntimeError('Unreachable code') -async def wait_for_all_tasks_for_finish( +async def wait_for_all_tasks_to_finish( tasks: Sequence[asyncio.Task], *, logger: Logger, diff --git a/src/crawlee/events/_event_manager.py b/src/crawlee/events/_event_manager.py index e43add30e9..49bf6d65c7 100644 --- a/src/crawlee/events/_event_manager.py +++ b/src/crawlee/events/_event_manager.py @@ -6,16 +6,13 @@ import inspect from collections import defaultdict from datetime import timedelta -from functools import wraps from logging import getLogger from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast, overload -from pyee.asyncio import AsyncIOEventEmitter - from crawlee._utils.context import ensure_context from crawlee._utils.docs import docs_group from crawlee._utils.recurring_task import RecurringTask -from crawlee._utils.wait import wait_for_all_tasks_for_finish +from crawlee._utils.wait import wait_for_all_tasks_to_finish from crawlee.events._types import ( Event, EventAbortingData, @@ -31,7 +28,7 @@ from collections.abc import Awaitable, Callable from types import TracebackType - from typing_extensions import NotRequired + from typing_extensions import NotRequired, Self from crawlee.events._types import EventData, WrappedListener @@ -55,9 +52,9 @@ class EventManagerOptions(TypedDict): class EventManager: """Manage events and their listeners, enabling registration, emission, and execution control. - It allows for registering event listeners, emitting events, and ensuring all listeners complete their execution. - Built on top of `pyee.asyncio.AsyncIOEventEmitter`. It implements additional features such as waiting for all - listeners to complete and emitting `PersistState` events at regular intervals. + Listeners can be registered for any of the events and are invoked whenever the event is emitted, each of them + in its own task - the sync ones in a separate thread so that they cannot block the event loop. On top of that, + it emits `PersistState` events at regular intervals and can wait for all the running listeners to complete. """ def __init__( @@ -73,77 +70,80 @@ def __init__( close_timeout: Optional timeout for canceling pending event listeners if they exceed this duration. """ self._persist_state_interval = persist_state_interval + """Interval between the emitted `PersistState` events.""" + self._close_timeout = close_timeout + """Timeout for the pending listeners when the context is being left, unlimited if not set.""" - # Asynchronous event emitter for handle events and invoke the event listeners. - self._event_emitter = AsyncIOEventEmitter() + self._listener_tasks: set[asyncio.Task[None]] = set() + """Tasks of the running listener invocations, kept both to keep them alive and to be able to wait for them.""" - # Listeners are wrapped inside asyncio.Task. Store their references here so that we can wait for them to finish. - self._listener_tasks: set[asyncio.Task] = set() + self._waiting_listener_tasks: set[asyncio.Task[None]] = set() + """Listener tasks currently blocked in `wait_for_all_listeners_to_complete`. - # Tasks currently blocked in `wait_for_all_listeners_to_complete`; excluded when gathering to avoid deadlock. - self._waiting_listener_tasks: set[asyncio.Task] = set() + They are excluded from the waits made by the listeners themselves, as a waiter must never await itself. + """ - # Store the mapping between events, listeners and their wrappers in the following way: - # event -> listener -> [wrapped_listener_1, wrapped_listener_2, ...] self._listeners_to_wrappers: dict[Event, dict[EventListener[Any], list[WrappedListener]]] = defaultdict( lambda: defaultdict(list), ) + """Registered listeners and their wrappers, mapped as `event -> listener -> [wrapper_1, wrapper_2, ...]`.""" - # Recurring task for emitting persist state events. self._emit_persist_state_event_rec_task = RecurringTask( func=self._emit_persist_state_event, delay=self._persist_state_interval, ) + """Recurring task emitting the `PersistState` events.""" - # Reference count for active contexts. self._active_ref_count = 0 + """Reference count of the active contexts.""" @property def active(self) -> bool: """Indicate whether the context is active.""" return self._active_ref_count > 0 - async def __aenter__(self) -> EventManager: + async def __aenter__(self) -> Self: """Initialize the event manager upon entering the async context.""" self._active_ref_count += 1 - if self._active_ref_count > 1: - return self - self._emit_persist_state_event_rec_task.start() + if self._active_ref_count == 1: + self._emit_persist_state_event_rec_task.start() + return self + @ensure_context async def __aexit__( self, exc_type: type[BaseException] | None, exc_value: BaseException | None, exc_traceback: TracebackType | None, ) -> None: - """Close the local event manager upon exiting the async context. + """Close the event manager upon exiting the async context. This will stop listening for the events, and it will wait for all the event listeners to finish. Raises: RuntimeError: If the context manager is not active. """ - if not self.active: - raise RuntimeError(f'The {self.__class__.__name__} is not active.') + last_exit = self._active_ref_count == 1 + + try: + # Stop the periodic emission first, so that the event emitted below really is the last one. + if last_exit: + await self._emit_persist_state_event_rec_task.stop() - if self._active_ref_count > 1: - # Emit persist state event to ensure the latest state is saved before closing the context. + # Emit a persist state event to ensure the latest state is saved before closing the context. await self._emit_persist_state_event() await self.wait_for_all_listeners_to_complete(timeout=self._close_timeout) - self._active_ref_count -= 1 - return + finally: + # The context has to be left even if closing fails, be it a cancellation or a failed emission. Staying + # active would mean never emitting `PersistState` again, as re-entering the context would be a no-op. + if last_exit: + self._listener_tasks.clear() + self._listeners_to_wrappers.clear() - # Stop persist state event periodic emission and manually emit last one to ensure latest state is saved. - await self._emit_persist_state_event_rec_task.stop() - await self._emit_persist_state_event() - await self.wait_for_all_listeners_to_complete(timeout=self._close_timeout) - self._event_emitter.remove_all_listeners() - self._listener_tasks.clear() - self._listeners_to_wrappers.clear() - self._active_ref_count -= 1 + self._active_ref_count -= 1 @overload def on(self, *, event: Literal[Event.PERSIST_STATE], listener: EventListener[EventPersistStateData]) -> None: ... @@ -167,49 +167,7 @@ def on(self, *, event: Event, listener: EventListener[Any]) -> None: event: The event for which to listen to. listener: The function (sync or async) which is to be called when the event is emitted. """ - signature = inspect.signature(listener) - - @wraps(cast('Callable[..., Awaitable[None] | None]', listener)) - async def listener_wrapper(event_data: EventData) -> None: - try: - bound_args = signature.bind(event_data) - except TypeError: # Parameterless listener - bound_args = signature.bind() - - # If the listener is a coroutine function, just call it, otherwise, run it in a separate thread - # to avoid blocking the event loop - coro = ( - listener(*bound_args.args, **bound_args.kwargs) - if inspect.iscoroutinefunction(listener) - else asyncio.to_thread(cast('Callable[..., None]', listener), *bound_args.args, **bound_args.kwargs) - ) - - listener_name = listener.__name__ if hasattr(listener, '__name__') else listener.__class__.__name__ - listener_task = asyncio.create_task(coro, name=f'Task-{event.value}-{listener_name}') - self._listener_tasks.add(listener_task) - - try: - logger.debug('EventManager.on.listener_wrapper(): Awaiting listener task...') - await listener_task - logger.debug('EventManager.on.listener_wrapper(): Listener task completed.') - except Exception: - # We need to swallow the exception and just log it here, otherwise it could break the event emitter - logger.exception( - 'Exception in the event listener', - extra={ - 'event_name': event.value, - 'listener_name': listener.__name__ - if hasattr(listener, '__name__') - else listener.__class__.__name__, - }, - ) - finally: - logger.debug('EventManager.on.listener_wrapper(): Removing listener task from the set...') - # `discard`, not `remove`: `__aexit__` may have cleared the set while this listener ran. - self._listener_tasks.discard(listener_task) - - self._listeners_to_wrappers[event][listener].append(listener_wrapper) - self._event_emitter.add_listener(event.value, listener_wrapper) + self._listeners_to_wrappers[event][listener].append(self._wrap_listener(event, listener)) def off(self, *, event: Event, listener: EventListener[Any] | None = None) -> None: """Remove a specific listener or all listeners for an event. @@ -219,13 +177,14 @@ def off(self, *, event: Event, listener: EventListener[Any] | None = None) -> No listener: The listener which is supposed to be removed. If not passed, all listeners of this event are removed. """ - if listener: - for listener_wrapper in self._listeners_to_wrappers[event][listener]: - self._event_emitter.remove_listener(event.value, listener_wrapper) - self._listeners_to_wrappers[event][listener] = [] - else: - self._listeners_to_wrappers[event] = defaultdict(list) - self._event_emitter.remove_all_listeners(event.value) + # Explicit `None` check - a listener may be an object that is falsy in a boolean context. + if listener is None: + self._listeners_to_wrappers.pop(event, None) + return + + # Popping the whole entry - an empty list of wrappers would keep a reference to a listener that is no longer + # registered, and looking the mapping up with `get` does not create entries for unknown events. + self._listeners_to_wrappers.get(event, {}).pop(listener, None) @overload def emit(self, *, event: Literal[Event.PERSIST_STATE], event_data: EventPersistStateData) -> None: ... @@ -243,14 +202,25 @@ def emit(self, *, event: Literal[Event.CRAWLER_STATUS], event_data: EventCrawler def emit(self, *, event: Event, event_data: Any) -> None: ... @ensure_context - def emit(self, *, event: Event, event_data: EventData) -> None: + def emit(self, *, event: Event, event_data: Any) -> None: """Emit an event with the associated data to all registered listeners. + Each listener is invoked in its own task, so this method only starts them. Use + `wait_for_all_listeners_to_complete` to wait for them to finish. + Args: event: The event which will be emitted. event_data: The data which will be passed to the event listeners. """ - self._event_emitter.emit(event.value, event_data) + # No listener can run before this method returns, as it does not await anything. Reading the mapping with + # `get` keeps the events that nobody listens to out of it. + for listener, listener_wrappers in self._listeners_to_wrappers.get(event, {}).items(): + task_name = f'Task-{event.value}-{self._get_listener_name(listener)}' + + for listener_wrapper in listener_wrappers: + listener_task = asyncio.create_task(listener_wrapper(event_data), name=task_name) + self._listener_tasks.add(listener_task) + listener_task.add_done_callback(self._listener_tasks.discard) @ensure_context async def wait_for_all_listeners_to_complete(self, *, timeout: timedelta | None = None) -> None: @@ -263,25 +233,78 @@ async def wait_for_all_listeners_to_complete(self, *, timeout: timedelta | None # A waiter can't finish until the listeners it awaits do, so waiters must never await each other or # themselves - this is what happens when a listener waits or closes from within itself. Only a waiter # that is a listener itself can be awaited this way, so only such waiters are tracked and excluded. - waiting_task = asyncio.current_task() - is_listener_waiter = waiting_task is not None and waiting_task in self._listener_tasks - if is_listener_waiter and waiting_task is not None: - self._waiting_listener_tasks.add(waiting_task) - - # `emit` only schedules the listener wrappers; each registers its listener task once it starts running, - # so yield first - otherwise listeners of a just-emitted event are missed and the wait is a no-op. - await asyncio.sleep(0) + current_task = asyncio.current_task() + listener_waiter = current_task if current_task in self._listener_tasks else None - # Any other caller is outside the cycle, so it must await every listener, waiting ones included. - excluded = self._waiting_listener_tasks if is_listener_waiter else frozenset[asyncio.Task]() - listener_tasks = [task for task in self._listener_tasks if task not in excluded] + if listener_waiter is not None: + self._waiting_listener_tasks.add(listener_waiter) + listener_tasks = [task for task in self._listener_tasks if task not in self._waiting_listener_tasks] + else: + # Any other caller is outside the cycle, so it must await every listener, the waiting ones included. + listener_tasks = list(self._listener_tasks) try: - await wait_for_all_tasks_for_finish(tasks=listener_tasks, logger=logger, timeout=timeout) + await wait_for_all_tasks_to_finish(tasks=listener_tasks, logger=logger, timeout=timeout) finally: - if is_listener_waiter and waiting_task is not None: - self._waiting_listener_tasks.discard(waiting_task) + if listener_waiter is not None: + self._waiting_listener_tasks.discard(listener_waiter) async def _emit_persist_state_event(self) -> None: """Emit a persist state event with the given migration status.""" self.emit(event=Event.PERSIST_STATE, event_data=EventPersistStateData(is_migrating=False)) + + @staticmethod + def _wrap_listener(event: Event, listener: EventListener[Any]) -> WrappedListener: + """Wrap a listener into a coroutine function that invokes it with the emitted event data. + + Whatever depends only on the listener itself is resolved here, once per registration, instead of on every + invocation. A sync listener is invoked in a separate thread so that it cannot block the event loop, and every + exception the invocation raises is logged and swallowed - one failing listener must not affect the others. + """ + is_async = EventManager._is_async_listener(listener) + takes_event_data = EventManager._takes_event_data(listener) + listener_name = EventManager._get_listener_name(listener) + + async def listener_wrapper(event_data: EventData) -> None: + args = (event_data,) if takes_event_data else () + + try: + if is_async: + await cast('Callable[..., Awaitable[None]]', listener)(*args) + else: + await asyncio.to_thread(cast('Callable[..., None]', listener), *args) + except Exception: + logger.exception( + 'Exception in the event listener', + extra={'event_name': event.value, 'listener_name': listener_name}, + ) + + return listener_wrapper + + @staticmethod + def _is_async_listener(listener: EventListener[Any]) -> bool: + """Check whether calling the listener returns a coroutine that has to be awaited. + + A plain `inspect.iscoroutinefunction` call sees only the object itself, so it reports a class instance with + an `async def __call__` - a perfectly valid listener - as synchronous. Its `__call__` has to be checked too. + """ + return inspect.iscoroutinefunction(listener) or inspect.iscoroutinefunction(listener.__call__) + + @staticmethod + def _takes_event_data(listener: EventListener[Any]) -> bool: + """Check whether the listener accepts the emitted event data, as a parameterless listener is allowed as well. + + Binding checks the parameters of the signature, not the value bound to them, so anything can stand in for + the event data here. + """ + try: + inspect.signature(listener).bind(None) + except TypeError: + return False + + return True + + @staticmethod + def _get_listener_name(listener: EventListener[Any]) -> str: + """Get a name of the listener for logging and for naming the tasks it runs in.""" + return getattr(listener, '__name__', None) or type(listener).__name__ diff --git a/src/crawlee/events/_local_event_manager.py b/src/crawlee/events/_local_event_manager.py index fbb08daab7..97b92150c5 100644 --- a/src/crawlee/events/_local_event_manager.py +++ b/src/crawlee/events/_local_event_manager.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: from types import TracebackType - from typing_extensions import Unpack + from typing_extensions import Self, Unpack logger = getLogger(__name__) @@ -45,15 +45,16 @@ def __init__( system_info_interval: Interval at which `SystemInfo` events are emitted. event_manager_options: Additional options for the parent class. """ + super().__init__(**event_manager_options) + self._system_info_interval = system_info_interval + """Interval between the emitted `SystemInfo` events.""" - # Recurring task for emitting system info events. self._emit_system_info_event_rec_task = RecurringTask( func=self._emit_system_info_event, delay=self._system_info_interval, ) - - super().__init__(**event_manager_options) + """Recurring task emitting the `SystemInfo` events.""" @classmethod def from_config(cls, config: Configuration | None = None) -> LocalEventManager: @@ -69,7 +70,7 @@ def from_config(cls, config: Configuration | None = None) -> LocalEventManager: persist_state_interval=config.persist_state_interval, ) - async def __aenter__(self) -> LocalEventManager: + async def __aenter__(self) -> Self: """Initialize the local event manager upon entering the async context. It starts emitting system info events at regular intervals. @@ -98,8 +99,12 @@ async def __aexit__( async def _emit_system_info_event(self) -> None: """Emit a system info event with the current CPU and memory usage.""" - cpu_info = await asyncio.to_thread(get_cpu_info) - memory_info = await asyncio.to_thread(get_memory_info) + # Both readings block the thread they run in - `get_cpu_info` even samples the CPU utilization over a short + # interval - so run them concurrently instead of one after the other. + cpu_info, memory_info = await asyncio.gather( + asyncio.to_thread(get_cpu_info), + asyncio.to_thread(get_memory_info), + ) event_data = EventSystemInfoData(cpu_info=cpu_info, memory_info=memory_info) self.emit(event=Event.SYSTEM_INFO, event_data=event_data) diff --git a/src/crawlee/storages/_request_queue.py b/src/crawlee/storages/_request_queue.py index e29bda4af9..b614928a2b 100644 --- a/src/crawlee/storages/_request_queue.py +++ b/src/crawlee/storages/_request_queue.py @@ -9,7 +9,7 @@ from crawlee import Request, service_locator from crawlee._utils.docs import docs_group -from crawlee._utils.wait import wait_for_all_tasks_for_finish +from crawlee._utils.wait import wait_for_all_tasks_to_finish from crawlee.request_loaders import RequestManager from crawlee.storage_clients.models import AddRequestsResponse @@ -242,7 +242,7 @@ async def _process_remaining_batches() -> None: # Wait for all tasks to finish if requested if wait_for_all_requests_to_be_added: - await wait_for_all_tasks_for_finish( + await wait_for_all_tasks_to_finish( (remaining_batches_task,), logger=logger, timeout=wait_for_all_requests_to_be_added_timeout, diff --git a/tests/unit/events/test_event_manager.py b/tests/unit/events/test_event_manager.py index aac81bd11f..772b554c1b 100644 --- a/tests/unit/events/test_event_manager.py +++ b/tests/unit/events/test_event_manager.py @@ -2,6 +2,7 @@ import asyncio import logging +import threading from contextlib import suppress from datetime import timedelta from functools import update_wrapper @@ -11,15 +12,18 @@ import pytest +from crawlee._utils.wait import wait_for_all_tasks_to_finish from crawlee.events import Event, EventManager, EventSystemInfoData if TYPE_CHECKING: - from collections.abc import AsyncGenerator + from collections.abc import AsyncGenerator, Sequence @pytest.fixture async def event_manager() -> AsyncGenerator[EventManager, None]: - async with EventManager() as event_manager: + # The teardown waits for the listeners left running by the test, so cap it - without a timeout, a listener that + # never completes hangs the whole run instead of failing the test that left it behind. + async with EventManager(close_timeout=timedelta(seconds=5)) as event_manager: yield event_manager @@ -48,6 +52,34 @@ def sync_listener(payload: Any) -> None: return sl +class ListenerWaitSpy: + """Observes the waits the event manager makes, patched in over the real waiting helper on construction. + + It lets a test see which tasks a wait awaits and act at the exact moment a wait starts, so that ordering does + not have to be approximated by sleeping. + """ + + def __init__(self, monkeypatch: pytest.MonkeyPatch) -> None: + self.awaited_tasks: list[set[asyncio.Task[None]]] = [] + """Tasks awaited by each of the waits, in the order the waits started.""" + + self.wait_started = asyncio.Event() + """Set by a wait right before it starts awaiting, so once it is observed, the waiter is already blocked.""" + + monkeypatch.setattr('crawlee.events._event_manager.wait_for_all_tasks_to_finish', self.wait) + + async def wait( + self, + tasks: Sequence[asyncio.Task[None]], + *, + logger: logging.Logger, + timeout: timedelta | None = None, + ) -> None: + self.awaited_tasks.append(set(tasks)) + self.wait_started.set() + await wait_for_all_tasks_to_finish(tasks, logger=logger, timeout=timeout) + + async def test_emit_invokes_registered_sync_listener( sync_listener: MagicMock, event_manager: EventManager, @@ -81,6 +113,28 @@ async def test_emit_invokes_both_sync_and_async_listeners( assert sync_listener.call_args[0] == (event_system_info_data,) +async def test_emit_starts_a_task_for_every_listener_invocation( + sync_listener: MagicMock, + async_listener: AsyncMock, + event_manager: EventManager, + event_system_info_data: EventSystemInfoData, +) -> None: + """Listener tasks are registered by `emit` itself, so a wait that follows it cannot miss any of them.""" + event_manager.on(event=Event.SYSTEM_INFO, listener=sync_listener) + # The same listener registered twice is invoked twice. + event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) + event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) + + event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + + assert len(event_manager._listener_tasks) == 3 + + await event_manager.wait_for_all_listeners_to_complete() + + assert sync_listener.call_count == 1 + assert async_listener.call_count == 2 + + async def test_emit_event_with_no_listeners( event_manager: EventManager, event_system_info_data: EventSystemInfoData, @@ -121,6 +175,64 @@ async def async_listener() -> None: assert async_mock.call_count == 1 +async def test_emit_logs_a_listener_that_cannot_be_called( + event_manager: EventManager, + event_system_info_data: EventSystemInfoData, + caplog: pytest.LogCaptureFixture, +) -> None: + """A listener that fits neither call shape must be reported like any other failing one, not escape its task.""" + + async def two_parameter_listener(_first: Any, _second: Any) -> None: + pass + + event_manager.on(event=Event.SYSTEM_INFO, listener=two_parameter_listener) # ty: ignore[no-matching-overload] + event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + + with caplog.at_level(logging.ERROR): + await event_manager.wait_for_all_listeners_to_complete() + + assert [record.message for record in caplog.records] == ['Exception in the event listener'] + + +async def test_emit_invokes_async_callable_instance_listener( + event_manager: EventManager, + event_system_info_data: EventSystemInfoData, +) -> None: + """A class instance with an `async def __call__` is an async listener - it must be awaited, not run in a thread.""" + received: list[Any] = [] + + class AsyncCallableListener: + async def __call__(self, event_data: Any) -> None: + received.append(event_data) + + event_manager.on(event=Event.SYSTEM_INFO, listener=AsyncCallableListener()) + event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + await event_manager.wait_for_all_listeners_to_complete() + + assert received == [event_system_info_data] + + +async def test_emit_invokes_sync_callable_instance_listener( + event_manager: EventManager, + event_system_info_data: EventSystemInfoData, +) -> None: + """A class instance with a plain `__call__` is a sync listener and is run in a separate thread.""" + received: list[Any] = [] + listener_thread_ids: list[int] = [] + + class SyncCallableListener: + def __call__(self, event_data: Any) -> None: + received.append(event_data) + listener_thread_ids.append(threading.get_ident()) + + event_manager.on(event=Event.SYSTEM_INFO, listener=SyncCallableListener()) + event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + await event_manager.wait_for_all_listeners_to_complete() + + assert received == [event_system_info_data] + assert listener_thread_ids != [threading.get_ident()] + + async def test_remove_nonexistent_listener_does_not_fail( async_listener: AsyncMock, event_manager: EventManager, @@ -144,6 +256,43 @@ async def test_removed_listener_not_invoked_on_emit( assert async_listener.call_count == 0 +async def test_off_with_a_falsy_listener_removes_only_that_listener( + sync_listener: MagicMock, + event_manager: EventManager, + event_system_info_data: EventSystemInfoData, +) -> None: + """A listener that is falsy in a boolean context must not be mistaken for "no listener given".""" + falsy_listener = MagicMock() + falsy_listener.__bool__.return_value = False + + event_manager.on(event=Event.SYSTEM_INFO, listener=falsy_listener) + event_manager.on(event=Event.SYSTEM_INFO, listener=sync_listener) + event_manager.off(event=Event.SYSTEM_INFO, listener=falsy_listener) + event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + + await event_manager.wait_for_all_listeners_to_complete() + + assert falsy_listener.call_count == 0 + assert sync_listener.call_count == 1 + + +async def test_off_leaves_no_entries_behind( + async_listener: AsyncMock, + event_manager: EventManager, +) -> None: + """Removing a listener must not leave an entry holding a reference to it, nor register unknown events.""" + event_manager.off(event=Event.SYSTEM_INFO, listener=async_listener) + event_manager.off(event=Event.ABORTING) + + assert event_manager._listeners_to_wrappers == {} + + event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) + event_manager.off(event=Event.SYSTEM_INFO, listener=async_listener) + + # Reading with `get` - a plain lookup would create the very entry that is asserted to be gone. + assert event_manager._listeners_to_wrappers.get(Event.SYSTEM_INFO) == {} + + async def test_close_clears_listeners_and_tasks(async_listener: AsyncMock) -> None: async with EventManager() as event_manager: event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) @@ -169,26 +318,78 @@ async def test_close_after_emit_processes_event( assert len(event_manager._listeners_to_wrappers) == 0 +async def test_nested_context_tears_down_on_the_last_exit_only( + async_listener: AsyncMock, + event_system_info_data: EventSystemInfoData, +) -> None: + """A nested exit only flushes the running listeners, the outermost one tears the event manager down.""" + event_manager = EventManager() + + async with event_manager: + event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) + + async with event_manager: + event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + + assert async_listener.call_count == 1 + assert event_manager.active is True + assert len(event_manager._listeners_to_wrappers) == 1 + + event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + + assert async_listener.call_count == 2 + assert event_manager.active is False + assert len(event_manager._listeners_to_wrappers) == 0 + + +async def test_context_is_left_even_if_closing_fails(monkeypatch: pytest.MonkeyPatch) -> None: + """A manager left active by a failed close would never emit `PersistState` again - re-entering it is a no-op.""" + event_manager = EventManager() + await event_manager.__aenter__() + + async def raise_error() -> None: + raise RuntimeError('Emitting failed.') + + monkeypatch.setattr(event_manager, '_emit_persist_state_event', raise_error) + + with pytest.raises(RuntimeError, match=r'Emitting failed\.'): + await event_manager.__aexit__(None, None, None) + + assert event_manager.active is False + + async def test_wait_for_all_listeners_cancelled_error( monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, event_system_info_data: EventSystemInfoData, ) -> None: - # Simulate long-running listener tasks - async def long_running_listener() -> None: - await asyncio.sleep(10) - - # Define a side effect function that raises CancelledError - async def mock_async_wait(*_: Any, **__: Any) -> None: + """A cancelled wait must cancel the listeners it was awaiting and let the cancellation propagate.""" + never_set = asyncio.Event() + listener_cancelled = asyncio.Event() + + async def never_ending_listener() -> None: + try: + await never_set.wait() + except asyncio.CancelledError: + listener_cancelled.set() + raise + + async def cancel_the_wait(*_: Any, **__: Any) -> None: raise asyncio.CancelledError - with pytest.raises(asyncio.CancelledError), caplog.at_level(logging.WARNING): # noqa: PT012 - async with EventManager(close_timeout=timedelta(milliseconds=10)) as event_manager: - event_manager.on(event=Event.SYSTEM_INFO, listener=long_running_listener) - event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + event_manager = EventManager() + await event_manager.__aenter__() + event_manager.on(event=Event.SYSTEM_INFO, listener=never_ending_listener) + event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + + # Only `asyncio.wait` is replaced, so the cancellation is still handled by the real waiting code - which is what + # has to cancel the listener that never finishes on its own. + monkeypatch.setattr('asyncio.wait', cancel_the_wait) + + # Capped, as a listener left uncancelled would be awaited forever - the close has no timeout of its own here. + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(event_manager.__aexit__(None, None, None), timeout=5) - # Use monkeypatch to replace asyncio.wait with mock_async_wait - monkeypatch.setattr('asyncio.wait', mock_async_wait) + assert listener_cancelled.is_set() async def test_methods_raise_error_when_not_active(event_system_info_data: EventSystemInfoData) -> None: @@ -202,6 +403,9 @@ async def test_methods_raise_error_when_not_active(event_system_info_data: Event with pytest.raises(RuntimeError, match=r'EventManager is not active.'): await event_manager.wait_for_all_listeners_to_complete() + with pytest.raises(RuntimeError, match=r'EventManager is not active.'): + await event_manager.__aexit__(None, None, None) + async with event_manager: event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) await event_manager.wait_for_all_listeners_to_complete() @@ -214,16 +418,20 @@ async def test_wait_for_all_listeners_from_within_a_listener_does_not_deadlock( event_system_info_data: EventSystemInfoData, ) -> None: """Waiting from within a listener must not self-await, yet must still await the other listeners.""" + parked = asyncio.Event() + release_other_listener = asyncio.Event() other_listener_done = asyncio.Event() waiter_done = asyncio.Event() other_done_when_wait_returned: bool | None = None async def other_listener(_: Any) -> None: - await asyncio.sleep(0.2) + await release_other_listener.wait() other_listener_done.set() async def waiting_listener(_: Any) -> None: nonlocal other_done_when_wait_returned + # Set right before the wait, which registers the waiter and captures the tasks without yielding in between. + parked.set() await event_manager.wait_for_all_listeners_to_complete() other_done_when_wait_returned = other_listener_done.is_set() waiter_done.set() @@ -232,11 +440,14 @@ async def waiting_listener(_: Any) -> None: event_manager.on(event=Event.SYSTEM_INFO, listener=waiting_listener) event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) + # The other listener finishes only once the waiter is already blocked, so the wait cannot pass it by. + await asyncio.wait_for(parked.wait(), timeout=5) + release_other_listener.set() + await asyncio.wait_for(waiter_done.wait(), timeout=5) # No self-await deadlock, and the wait must have blocked until the co-registered listener finished. assert other_done_when_wait_returned is True - assert other_listener_done.is_set() async def test_wait_from_within_multiple_listeners_does_not_deadlock( @@ -266,81 +477,93 @@ async def second_waiting_listener(_: Any) -> None: async def test_wait_from_outside_awaits_a_listener_that_is_itself_waiting( + monkeypatch: pytest.MonkeyPatch, event_manager: EventManager, event_system_info_data: EventSystemInfoData, ) -> None: """A caller that is not a listener is outside the deadlock cycle, so it must await even waiting listeners.""" + spy = ListenerWaitSpy(monkeypatch) parked = asyncio.Event() + release_other_listener = asyncio.Event() waiting_listener_done = asyncio.Event() + waiter_task: asyncio.Task[None] | None = None async def other_listener(_: Any) -> None: - await asyncio.sleep(0.1) + await release_other_listener.wait() async def waiting_listener(_: Any) -> None: + nonlocal waiter_task + waiter_task = asyncio.current_task() + # Set right before the wait, which registers the waiter and captures the tasks without yielding in between. parked.set() await event_manager.wait_for_all_listeners_to_complete() - # Work done after the inner wait returns - the outer wait must not return before it finishes. - await asyncio.sleep(0.1) waiting_listener_done.set() event_manager.on(event=Event.SYSTEM_INFO, listener=other_listener) event_manager.on(event=Event.SYSTEM_INFO, listener=waiting_listener) event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) - # `parked` is set right before the listener registers itself as a waiter, without yielding in between. await asyncio.wait_for(parked.wait(), timeout=5) + # The release has to come from outside this task, as the wait below blocks until the listeners are done. Scheduled + # on the loop, it runs only once that wait has suspended - so the wait captures the tasks with both still running. + asyncio.get_running_loop().call_soon(release_other_listener.set) await asyncio.wait_for(event_manager.wait_for_all_listeners_to_complete(), timeout=5) + # The waiter's own wait went first, so the second one is the outer wait - and it had to include the parked waiter. + assert waiter_task in spy.awaited_tasks[1] assert waiting_listener_done.is_set() async def test_close_from_within_a_listener_does_not_deadlock_or_error( + monkeypatch: pytest.MonkeyPatch, event_system_info_data: EventSystemInfoData, ) -> None: """Closing the event manager from within a listener (as `Actor.exit()` does) must not deadlock or raise.""" + spy = ListenerWaitSpy(monkeypatch) event_manager = EventManager() await event_manager.__aenter__() - # A wrapper finalizing after close raises onto the loop (its `error` listener is gone by then), so watch - # both channels for a stray exception. - emitter_errors: list[BaseException] = [] - event_manager._event_emitter.add_listener('error', emitter_errors.append) + # A listener task finalizing after the close cleared the task set must not report a stray exception. loop_errors: list[dict[str, Any]] = [] asyncio.get_running_loop().set_exception_handler(lambda _loop, context: loop_errors.append(context)) closed = asyncio.Event() + release_other_listener = asyncio.Event() other_listener_done = asyncio.Event() + closing_task: asyncio.Task[None] | None = None async def other_listener(_: Any) -> None: - await asyncio.sleep(0.2) + await release_other_listener.wait() other_listener_done.set() async def closing_listener(_: Any) -> None: + nonlocal closing_task + closing_task = asyncio.current_task() await event_manager.__aexit__(None, None, None) closed.set() # A second listener makes close await a concurrently-running listener - the real `Actor.exit()` shape. event_manager.on(event=Event.SYSTEM_INFO, listener=other_listener) event_manager.on(event=Event.SYSTEM_INFO, listener=closing_listener) - - tasks_before = asyncio.all_tasks() event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) try: + # The close is already blocked in its wait here, so releasing the other listener only now means the close + # really does have to await a listener that is still running. + await asyncio.wait_for(spy.wait_started.wait(), timeout=5) + release_other_listener.set() + await asyncio.wait_for(closed.wait(), timeout=5) - # Drain the wrapper tasks so their `finally` blocks run before we assert - no arbitrary sleep. - spawned = asyncio.all_tasks() - tasks_before - {asyncio.current_task()} - if spawned: - await asyncio.wait(spawned) + # Let the closing listener's own task finalize before asserting - its wrapper returns right after `closed`. + assert closing_task is not None + await asyncio.wait_for(closing_task, timeout=5) finally: # Cap the cleanup so a regressed deadlock surfaces the real failure instead of hanging. if event_manager.active: with suppress(Exception): await asyncio.wait_for(event_manager.__aexit__(None, None, None), timeout=5) - # With `discard` no wrapper raises on finalize; the `remove` regression would surface on one of these. - assert emitter_errors == [] assert loop_errors == [] assert other_listener_done.is_set() assert event_manager.active is False diff --git a/tests/unit/events/test_local_event_manager.py b/tests/unit/events/test_local_event_manager.py index f2e998ab18..b88bba372d 100644 --- a/tests/unit/events/test_local_event_manager.py +++ b/tests/unit/events/test_local_event_manager.py @@ -1,15 +1,26 @@ from __future__ import annotations import asyncio +import threading from datetime import timedelta -from typing import Any -from unittest.mock import AsyncMock +from typing import TYPE_CHECKING, Any +from unittest.mock import AsyncMock, MagicMock +from crawlee._utils.system import CpuInfo, MemoryInfo from crawlee.events import LocalEventManager from crawlee.events._types import Event, EventSystemInfoData +if TYPE_CHECKING: + import pytest + + +async def test_emit_system_info_event(monkeypatch: pytest.MonkeyPatch) -> None: + """The recurring task emits the first `SystemInfo` event as soon as it starts, without waiting for the interval.""" + # Both readings are replaced with instant ones - a real `get_cpu_info` samples the CPU utilization over 100 ms, + # and on a loaded runner it takes far longer than that. + monkeypatch.setattr('crawlee.events._local_event_manager.get_cpu_info', lambda: MagicMock(spec=CpuInfo)) + monkeypatch.setattr('crawlee.events._local_event_manager.get_memory_info', lambda: MagicMock(spec=MemoryInfo)) -async def test_emit_system_info_event() -> None: mocked_listener = AsyncMock() received = asyncio.Event() @@ -17,13 +28,44 @@ async def async_listener(payload: Any) -> None: await mocked_listener(payload) received.set() - system_info_interval = timedelta(milliseconds=50) - async with LocalEventManager(system_info_interval=system_info_interval) as event_manager: + # An interval this long means the event can only come from the immediate first run of the recurring task. + async with LocalEventManager(system_info_interval=timedelta(hours=1)) as event_manager: + # Registered before anything yields to the event loop, so the very first emission already reaches it. event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) - # Wait until the listener is invoked at least once. A generous timeout avoids flakiness on - # loaded CI runners, where `psutil.cpu_percent(interval=0.1)` in `asyncio.to_thread` can - # delay the first emission well beyond the configured interval. - await asyncio.wait_for(received.wait(), timeout=10) + await asyncio.wait_for(received.wait(), timeout=5) assert mocked_listener.call_count >= 1 assert isinstance(mocked_listener.call_args[0][0], EventSystemInfoData) + + +async def test_system_info_readings_run_concurrently(monkeypatch: pytest.MonkeyPatch) -> None: + """Both readings block their thread, so they have to run at the same time - the barrier clears only if they do.""" + # A party left waiting alone breaks the barrier, which fails the test instead of hanging it. + barrier = threading.Barrier(2, timeout=5) + + def get_cpu_info_at_barrier() -> Any: + barrier.wait() + return MagicMock(spec=CpuInfo) + + def get_memory_info_at_barrier() -> Any: + barrier.wait() + return MagicMock(spec=MemoryInfo) + + monkeypatch.setattr('crawlee.events._local_event_manager.get_cpu_info', get_cpu_info_at_barrier) + monkeypatch.setattr('crawlee.events._local_event_manager.get_memory_info', get_memory_info_at_barrier) + + received: list[EventSystemInfoData] = [] + + async def listener(event_data: EventSystemInfoData) -> None: + received.append(event_data) + + async with LocalEventManager() as event_manager: + # A recurring emission would pair up with the direct one below at the barrier, so stop it first. It is + # cancelled before it ever runs, as nothing has yielded to the event loop since it was started. + await event_manager._emit_system_info_event_rec_task.stop() + + event_manager.on(event=Event.SYSTEM_INFO, listener=listener) + await event_manager._emit_system_info_event() + await event_manager.wait_for_all_listeners_to_complete() + + assert len(received) == 1 diff --git a/uv.lock b/uv.lock index 48afe76798..f0482d2035 100644 --- a/uv.lock +++ b/uv.lock @@ -781,7 +781,6 @@ dependencies = [ { name = "psutil" }, { name = "pydantic" }, { name = "pydantic-settings" }, - { name = "pyee" }, { name = "tldextract" }, { name = "typing-extensions" }, { name = "yarl" }, @@ -968,7 +967,6 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.11.0" }, { name = "pydantic-ai-slim", extras = ["openai"], marker = "extra == 'pydantic-ai'", specifier = ">=2.1.0" }, { name = "pydantic-settings", specifier = ">=2.12.0" }, - { name = "pyee", specifier = ">=9.0.0" }, { name = "redis", extras = ["hiredis"], marker = "extra == 'redis'", specifier = ">=7.0.0" }, { name = "rich", marker = "extra == 'cli'", specifier = ">=13.9.0" }, { name = "scikit-learn", marker = "extra == 'adaptive-crawler'", specifier = ">=1.6.0" },