feat(crawlers): Add BasicCrawler.pause() and resume(), and pause on migration

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,

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,14 +71,32 @@ 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."""
         return self._stop_requested
 
     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.

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.
 

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'

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)