From e7ff66b394ec5fdd11b9b78e97f584198e550421 Mon Sep 17 00:00:00 2001 From: Max Bohomolov Date: Mon, 9 Nov 2026 23:38:02 +0000 Subject: [PATCH] add `BasicCrawler.pause()` and `resume()`, and pause on migration --- src/crawlee/crawlers/_basic/_basic_crawler.py | 59 ++++++- src/crawlee/crawlers/_basic/_crawler_run.py | 21 ++- src/crawlee/statistics/_statistics.py | 4 + tests/unit/_statistics/test_persistence.py | 14 ++ .../crawlers/_basic/test_basic_crawler.py | 153 +++++++++++++++++- 5 files changed, 248 insertions(+), 3 deletions(-) diff --git a/src/crawlee/crawlers/_basic/_basic_crawler.py b/src/crawlee/crawlers/_basic/_basic_crawler.py index 5865b5d369..63bc4d556e 100644 --- a/src/crawlee/crawlers/_basic/_basic_crawler.py +++ b/src/crawlee/crawlers/_basic/_basic_crawler.py @@ -284,6 +284,9 @@ class BasicCrawler(Generic[TCrawlingContext, TStatisticsState]): _request_handler_timeout_text = 'Request handler timed out after' __next_id = 0 + _SAFE_MIGRATION_WAIT = timedelta(seconds=20) + """How long a migration or an abort waits for the requests in progress.""" + def __init__( self, *, @@ -581,6 +584,11 @@ def _live_run(self) -> CrawlerRun[TStatisticsState] | None: """The run in progress, if any.""" return self._last_run if self._last_run is not None and self._last_run.is_live else None + @property + def _dispatching_run(self) -> CrawlerRun[TStatisticsState] | None: + """The run whose task loop is dispatching, if any. Only that one can be paused and resumed.""" + return self._live_run if self._live_run is not None and self._live_run.is_dispatching else None + def stop(self, reason: str = 'Stop was called externally.') -> None: """Set flag to stop crawler. @@ -598,6 +606,33 @@ def stop(self, reason: str = 'Stop was called externally.') -> None: # A caller can stop a run it scheduled with `create_task(crawler.run())` before the task starts. self._stop_before_run = True + async def pause(self, *, timeout: timedelta | None = None) -> None: + """Stop dispatching new requests and wait until the ones in progress finish. + + Unlike `stop`, this doesn't end the run: `run` stays pending until `resume` is called. Awaited from a request + handler, it waits for that handler's own request too. + + Args: + timeout: How long to wait for the requests in progress. The crawler stays paused if they don't finish in + time. + + Raises: + TimeoutError: If the requests in progress don't finish within `timeout`. + """ + if self._dispatching_run is None: + self._logger.warning('Cannot pause a crawler that is not running.') + return + + await self._dispatching_run.pause(timeout=timeout) + + def resume(self) -> None: + """Resume a crawler paused with `pause`, so that it dispatches requests again.""" + if self._dispatching_run is None: + self._logger.warning('Cannot resume a crawler that is not running.') + return + + self._dispatching_run.resume() + def _wrap_handler_with_error_context( self, handler: Callable[[TCrawlingContext | BasicCrawlingContext, Exception], Awaitable[T]] ) -> Callable[[TCrawlingContext | BasicCrawlingContext, Exception], Awaitable[T]]: @@ -922,7 +957,29 @@ async def _run_crawler(self) -> None: for context in contexts_to_enter: await exit_stack.enter_async_context(context) - await self._run.dispatch_requests() + event_manager = self._service_locator.get_event_manager() + event_manager.on(event=Event.MIGRATING, listener=self._pause_by_event) + event_manager.on(event=Event.ABORTING, listener=self._pause_by_event) + + try: + await self._run.dispatch_requests() + finally: + event_manager.off(event=Event.MIGRATING, listener=self._pause_by_event) + event_manager.off(event=Event.ABORTING, listener=self._pause_by_event) + + async def _pause_by_event(self) -> None: + """Pause the crawl before the process migrates or is aborted, so that no request is left half-done.""" + if self._dispatching_run is not None: + try: + await self._dispatching_run.pause(timeout=self._SAFE_MIGRATION_WAIT) + except TimeoutError: + self._logger.error( # noqa: TRY400 + 'The crawler was paused due to migration to another host, but some requests did not finish in ' + "time. Those requests' results may be duplicated." + ) + + # The persist state event that comes with a migration doesn't wait for the requests drained above. + await self._statistics.persist_state() async def add_requests( self, diff --git a/src/crawlee/crawlers/_basic/_crawler_run.py b/src/crawlee/crawlers/_basic/_crawler_run.py index ec602cc8b9..c02295ff97 100644 --- a/src/crawlee/crawlers/_basic/_crawler_run.py +++ b/src/crawlee/crawlers/_basic/_crawler_run.py @@ -52,6 +52,7 @@ def __init__( ) self._logger_once = LoggerOnce(logger) self._stop_requested = False + self._dispatching = False self._finished = False self._wake_at: datetime | None = None self._wake_timer: asyncio.TimerHandle | None = None @@ -70,6 +71,11 @@ def is_live(self) -> bool: """Whether the `BasicCrawler.run` call this belongs to is still executing.""" return not self._finished + @property + def is_dispatching(self) -> bool: + """Whether the task loop is dispatching requests. Pausing a run before then would keep it from ever starting.""" + return self._dispatching + @property def stop_requested(self) -> bool: """Whether a graceful shutdown of this run has been requested.""" @@ -77,7 +83,20 @@ def stop_requested(self) -> bool: async def dispatch_requests(self) -> None: """Run the task loop until the crawl is over. The concurrency system has to be active by now.""" - await self._autoscaled_pool.run() + self._dispatching = True + + try: + await self._autoscaled_pool.run() + finally: + self._dispatching = False + + async def pause(self, *, timeout: timedelta | None = None) -> None: + """Stop dispatching new requests and wait until the ones in progress finish.""" + await self._autoscaled_pool.pause(timeout=timeout) + + def resume(self) -> None: + """Resume dispatching after `pause`.""" + self._autoscaled_pool.resume() def schedule_wake(self, ready_at: datetime) -> None: """Wake the task loop at `ready_at`, unless an earlier wake-up is already scheduled. diff --git a/src/crawlee/statistics/_statistics.py b/src/crawlee/statistics/_statistics.py index 02bd4c7bee..74ff13e5ba 100644 --- a/src/crawlee/statistics/_statistics.py +++ b/src/crawlee/statistics/_statistics.py @@ -293,6 +293,10 @@ def calculate(self) -> FinalStatistics: retry_histogram=serialized_state['request_retry_histogram'], ) + async def persist_state(self) -> None: + """Persist the current statistics, if persistence is enabled.""" + await self._state.persist_state() + def reset(self) -> None: """Reset the statistics to their defaults, leaving the persisted record alone. diff --git a/tests/unit/_statistics/test_persistence.py b/tests/unit/_statistics/test_persistence.py index 574b57f382..cc00ef41d1 100644 --- a/tests/unit/_statistics/test_persistence.py +++ b/tests/unit/_statistics/test_persistence.py @@ -240,6 +240,20 @@ async def test_persistence_by_id(second_id: str, expected_requests_failed: int) assert statistics.state.requests_failed == expected_requests_failed +async def test_persist_state_writes_state_in_progress() -> None: + """`persist_state()` writes the current state while the statistics are still active.""" + key = 'statistics_persist_state' + kvs = await KeyValueStore.open() + + async with Statistics.with_default_state(persistence_enabled=True, persist_state_key=key) as statistics: + statistics.state.requests_failed = 42 + await statistics.persist_state() + persisted = await kvs.get_value(key) + + assert persisted is not None + assert persisted['requestsFailed'] == 42 + + async def test_persistence_loads_numeric_id() -> None: """A persisted state whose `statsId` is a number loads with a string ID.""" key = 'statistics_foo' diff --git a/tests/unit/crawlers/_basic/test_basic_crawler.py b/tests/unit/crawlers/_basic/test_basic_crawler.py index 093a7215c2..dca7753551 100644 --- a/tests/unit/crawlers/_basic/test_basic_crawler.py +++ b/tests/unit/crawlers/_basic/test_basic_crawler.py @@ -37,7 +37,13 @@ SessionError, UserDefinedErrorHandlerError, ) -from crawlee.events import Event, EventCrawlerStatusData, LocalEventManager +from crawlee.events import ( + Event, + EventAbortingData, + EventCrawlerStatusData, + EventMigratingData, + LocalEventManager, +) from crawlee.http_clients import HttpClient, HttpResponse from crawlee.request_loaders import ( FinishedStatus, @@ -1947,6 +1953,151 @@ async def test_run_while_running() -> None: await asyncio.wait_for(run_task, timeout=10) +async def test_pause_holds_requests_until_resume() -> None: + """`pause()` waits for the request in progress, and the paused run starts nothing new until `resume()`.""" + crawler = BasicCrawler(concurrency_settings=ConcurrencySettings(desired_concurrency=1, max_concurrency=1)) + release = asyncio.Event() + visited = list[str]() + + @crawler.router.default_handler + async def handler(context: BasicCrawlingContext) -> None: + visited.append(context.request.url) + await release.wait() + + urls = ['https://a.placeholder.com', 'https://b.placeholder.com'] + run_task = asyncio.create_task(crawler.run(urls)) + assert await poll_until_condition(lambda: len(visited) == 1) + + pause_task = asyncio.create_task(crawler.pause()) + await asyncio.sleep(0) + assert not pause_task.done() + + release.set() + await asyncio.wait_for(pause_task, timeout=10) + + # Give a pause that doesn't hold the task loop the chance to start the next request. + await asyncio.sleep(0.1) + assert visited == urls[:1] + assert not run_task.done() + + crawler.resume() + await asyncio.wait_for(run_task, timeout=10) + + assert visited == urls + + +async def test_pause_and_resume_before_run_warn(caplog: pytest.LogCaptureFixture) -> None: + """`pause()` and `resume()` before a run only log a warning.""" + crawler = BasicCrawler() + + with caplog.at_level(logging.WARNING, logger='crawlee'): + await crawler.pause() + crawler.resume() + + assert 'Cannot pause a crawler that is not running.' in caplog.text + assert 'Cannot resume a crawler that is not running.' in caplog.text + + +async def test_pause_and_resume_after_run_warn(caplog: pytest.LogCaptureFixture) -> None: + """`pause()` and `resume()` on a crawler whose run has finished only log a warning.""" + crawler = BasicCrawler(request_handler=AsyncMock()) + await crawler.run(['https://a.placeholder.com']) + + with caplog.at_level(logging.WARNING, logger='crawlee'): + await crawler.pause() + crawler.resume() + + assert 'Cannot pause a crawler that is not running.' in caplog.text + assert 'Cannot resume a crawler that is not running.' in caplog.text + + +@pytest.mark.parametrize( + ('event', 'event_data'), + [ + pytest.param(Event.MIGRATING, EventMigratingData(), id='migrating'), + pytest.param(Event.ABORTING, EventAbortingData(), id='aborting'), + ], +) +async def test_migration_pauses_crawl_and_persists_statistics( + event: Event, event_data: EventMigratingData | EventAbortingData +) -> None: + """A migration or an abort pauses the crawl and persists the statistics once the request in progress finishes.""" + crawler = BasicCrawler(concurrency_settings=ConcurrencySettings(desired_concurrency=1, max_concurrency=1)) + release = asyncio.Event() + visited = list[str]() + + @crawler.router.default_handler + async def handler(context: BasicCrawlingContext) -> None: + visited.append(context.request.url) + await release.wait() + + urls = ['https://a.placeholder.com', 'https://b.placeholder.com'] + run_task = asyncio.create_task(crawler.run(urls)) + assert await poll_until_condition(lambda: len(visited) == 1) + + finished_when_persisted = list[int]() + persist_state = AsyncMock( + side_effect=lambda: finished_when_persisted.append(crawler.statistics.state.requests_finished) + ) + event_manager = service_locator.get_event_manager() + + with patch.object(crawler.statistics, 'persist_state', persist_state): + event_manager.emit(event=event, event_data=event_data) + await asyncio.sleep(0) + persist_state.assert_not_called() + + release.set() + await event_manager.wait_for_all_listeners_to_complete() + + assert finished_when_persisted == [1] + assert visited == urls[:1] + + crawler.resume() + await asyncio.wait_for(run_task, timeout=10) + + +async def test_migration_timeout_logs_error(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None: + """A migration that the request in progress outlasts logs an error.""" + monkeypatch.setattr(BasicCrawler, '_SAFE_MIGRATION_WAIT', timedelta(milliseconds=50)) + crawler = BasicCrawler(concurrency_settings=ConcurrencySettings(desired_concurrency=1, max_concurrency=1)) + release = asyncio.Event() + visited = list[str]() + + @crawler.router.default_handler + async def handler(context: BasicCrawlingContext) -> None: + visited.append(context.request.url) + await release.wait() + + run_task = asyncio.create_task(crawler.run(['https://a.placeholder.com'])) + assert await poll_until_condition(lambda: len(visited) == 1) + + event_manager = service_locator.get_event_manager() + with caplog.at_level(logging.ERROR, logger='crawlee'): + event_manager.emit(event=Event.MIGRATING, event_data=EventMigratingData()) + await event_manager.wait_for_all_listeners_to_complete() + + assert 'some requests did not finish in time' in caplog.text + + release.set() + crawler.resume() + await asyncio.wait_for(run_task, timeout=10) + + +async def test_run_stops_listening_for_migration() -> None: + """Once `run()` returns, a migration no longer reaches the crawler.""" + event_manager = LocalEventManager() + crawler = BasicCrawler(event_manager=event_manager, request_handler=AsyncMock()) + + async with event_manager: + await crawler.run(['https://a.placeholder.com']) + + with patch.object(crawler.statistics, 'persist_state') as persist_state: + event_manager.emit(event=Event.MIGRATING, event_data=EventMigratingData()) + await event_manager.wait_for_all_listeners_to_complete() + + persist_state.assert_not_called() + + async def test_services_no_side_effect_on_crawler_init() -> None: custom_configuration = Configuration() custom_event_manager = LocalEventManager.from_config(custom_configuration)