feat(autoscaling)!: Make AutoscaledPool.pause() wait for the running tasks

docs/upgrading/upgrading_to_v2.md

@@ -155,6 +155,7 @@ The autoscaling classes now live in the public `crawlee.autoscaling` package, wh
 - `run_task_function` receives the id the task is booked under.
 - `current_concurrency` comes from the concurrency system, so a shared system may count the tasks of other pools too. The pool's own tasks are counted by the new `own_concurrency`.
 - `is_finished_function` is called only when no new task can start and none of the pool's own tasks are in flight, and never while the pool is paused. In v1, it was called on every loop iteration, so a paused pool could finish. A pool whose `is_task_ready_function` keeps resolving to `True` doesn't finish, so make it resolve to `False` once the work runs out.
+- `pause()` is a coroutine that returns once the pool's tasks in progress finish. With the new `timeout` option, it raises `TimeoutError` if they take longer, and the pool stays paused. Await it, since a call without `await` doesn't pause the pool.
 
 For example:
 

src/crawlee/autoscaling/_autoscaled_pool.py

@@ -145,10 +145,31 @@ async def abort(self) -> None:
         self._current_run.result.set_exception(AbortError())
         await self._current_run.cleanup_done.wait()
 
-    def pause(self) -> None:
-        """Pause the autoscaled pool so that it does not start new tasks."""
+    async def pause(self, *, timeout: timedelta | None = None) -> None:
+        """Stop starting new tasks and wait until the pool's tasks in progress finish.
+
+        Unlike `abort`, this doesn't end the run, which waits for `resume`. A call from inside `run_task_function` waits
+        for its own task too, so without a `timeout` it never returns.
+
+        Args:
+            timeout: How long to wait for the tasks in progress. The pool stays paused if they don't finish in time.
+
+        Raises:
+            TimeoutError: If the tasks in progress don't finish within `timeout`.
+        """
         self._is_paused = True
 
+        run = self._current_run
+        if run is None or not run.worker_tasks:
+            return
+
+        timeout_seconds = timeout.total_seconds() if timeout is not None else None
+        _, pending = await asyncio.wait(run.worker_tasks, timeout=timeout_seconds)
+        if pending:
+            raise TimeoutError(
+                f"The pool's running tasks did not finish within {timeout_seconds} seconds after `pause()` was called."
+            )
+
     def resume(self) -> None:
         """Resume a paused autoscaled pool so that it continues starting new tasks."""
         self._is_paused = False
@@ -203,6 +224,10 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None:
                     logger.debug('Not scheduling new task - no task is ready')
                     if finished := await self._maybe_finish(run):
                         break
+                elif self._is_paused:
+                    # The pool may have been paused while `is_task_ready_function` was awaited, and `pause` waits only
+                    # for the tasks that were already running.
+                    logger.debug('Not scheduling new tasks - the autoscaled pool is paused')
                 else:
                     task_id = str(self._next_task_id)
                     self._next_task_id += 1
@@ -251,9 +276,14 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None:
     async def _maybe_finish(self, run: _AutoscaledPoolRun) -> bool:
         """Ask `is_finished_function` whether the run is over, but only while none of this pool's tasks is in flight.
 
-        A running task can still produce new work, and with a shared system only this pool's own tasks count.
+        A running task can still produce new work, and with a shared system only this pool's own tasks count. A paused
+        pool never finishes, even when it was paused while `is_task_ready_function` or `is_finished_function` was
+        awaited.
         """
-        return not run.worker_tasks and await self._is_finished_function()
+        if self._is_paused or run.worker_tasks:
+            return False
+
+        return await self._is_finished_function() and not self._is_paused
 
     def _reap_worker_task(self, task: asyncio.Task, *, run: _AutoscaledPoolRun, task_id: str) -> None:
         """Handle cleanup and tracking of a completed worker task.

tests/unit/autoscaling/test_autoscaled_pool.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import logging
+from datetime import timedelta
 from itertools import chain, repeat
 from typing import TYPE_CHECKING, TypeVar
 from unittest.mock import AsyncMock, Mock, patch
@@ -192,7 +193,7 @@ async def test_does_not_finish_while_paused(make_concurrency_system: type[Concur
             is_task_ready_function=lambda: future(False),
             is_finished_function=is_finished,
         )
-        pool.pause()
+        await pool.pause()
 
         pool_run_task = asyncio.create_task(pool.run(), name='pool run task')
         try:
@@ -208,6 +209,238 @@ async def test_does_not_finish_while_paused(make_concurrency_system: type[Concur
     is_finished.assert_awaited_once()
 
 
+async def test_pause_waits_for_running_tasks(make_concurrency_system: type[ConcurrencySystem]) -> None:
+    """`pause` returns once the tasks in progress finish."""
+    release = asyncio.Event()
+
+    async def wait_for_release(_task_id: str) -> None:
+        await release.wait()
+
+    settings = ConcurrencySettings(min_concurrency=2, max_concurrency=2, desired_concurrency=2)
+
+    async with make_concurrency_system(settings) as concurrency_system:
+        pool = AutoscaledPool(
+            concurrency_system=concurrency_system,
+            consumer=CONSUMER,
+            run_task_function=AsyncMock(side_effect=wait_for_release),
+            is_task_ready_function=lambda: future(True),
+            is_finished_function=lambda: future(False),
+        )
+
+        pool_run_task = asyncio.create_task(pool.run(), name='pool run task')
+        try:
+            assert await poll_until_condition(lambda: pool.own_concurrency == 2)
+
+            pause_task = asyncio.create_task(pool.pause(), name='pause task')
+            await asyncio.sleep(0.1)
+            assert not pause_task.done()
+
+            release.set()
+            await asyncio.wait_for(pause_task, timeout=5)
+
+            assert pool.own_concurrency == 0
+        finally:
+            release.set()
+            pool_run_task.cancel()
+            await asyncio.gather(pool_run_task, return_exceptions=True)
+
+
+async def test_repeated_pause_waits(make_concurrency_system: type[ConcurrencySystem]) -> None:
+    """A `pause` on a pool that's already paused waits for the tasks in progress too."""
+    release = asyncio.Event()
+
+    async def wait_for_release(_task_id: str) -> None:
+        await release.wait()
+
+    settings = ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1)
+
+    async with make_concurrency_system(settings) as concurrency_system:
+        pool = AutoscaledPool(
+            concurrency_system=concurrency_system,
+            consumer=CONSUMER,
+            run_task_function=AsyncMock(side_effect=wait_for_release),
+            is_task_ready_function=lambda: future(True),
+            is_finished_function=lambda: future(False),
+        )
+
+        pool_run_task = asyncio.create_task(pool.run(), name='pool run task')
+        try:
+            assert await poll_until_condition(lambda: pool.own_concurrency == 1)
+
+            first_pause = asyncio.create_task(pool.pause(), name='first pause task')
+            second_pause = asyncio.create_task(pool.pause(), name='second pause task')
+            await asyncio.sleep(0.1)
+            assert not second_pause.done()
+
+            release.set()
+            await asyncio.wait_for(asyncio.gather(first_pause, second_pause), timeout=5)
+        finally:
+            release.set()
+            pool_run_task.cancel()
+            await asyncio.gather(pool_run_task, return_exceptions=True)
+
+
+async def test_pause_times_out(make_concurrency_system: type[ConcurrencySystem]) -> None:
+    """A `pause` that its tasks outlast raises `TimeoutError`, and the pool stays paused."""
+    release = asyncio.Event()
+
+    async def wait_for_release(_task_id: str) -> None:
+        await release.wait()
+
+    run = AsyncMock(side_effect=wait_for_release)
+    settings = ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1)
+
+    async with make_concurrency_system(settings) as concurrency_system:
+        pool = AutoscaledPool(
+            concurrency_system=concurrency_system,
+            consumer=CONSUMER,
+            run_task_function=run,
+            is_task_ready_function=lambda: future(True),
+            is_finished_function=lambda: future(False),
+        )
+
+        pool_run_task = asyncio.create_task(pool.run(), name='pool run task')
+        try:
+            assert await poll_until_condition(lambda: pool.own_concurrency == 1)
+
+            with pytest.raises(TimeoutError, match='did not finish'):
+                await pool.pause(timeout=timedelta(milliseconds=50))
+
+            release.set()
+            assert await poll_until_condition(lambda: pool.own_concurrency == 0)
+            await asyncio.sleep(0.1)
+            run.assert_awaited_once()
+        finally:
+            release.set()
+            pool_run_task.cancel()
+            await asyncio.gather(pool_run_task, return_exceptions=True)
+
+
+async def test_resume_during_pause(make_concurrency_system: type[ConcurrencySystem]) -> None:
+    """A `pause` overtaken by `resume` returns once the tasks it found finish, while newer tasks keep running."""
+    first_release = asyncio.Event()
+    later_release = asyncio.Event()
+
+    async def wait_for_release(task_id: str) -> None:
+        release = first_release if task_id == '0' else later_release
+        await release.wait()
+
+    run = AsyncMock(side_effect=wait_for_release)
+    settings = ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1)
+
+    async with make_concurrency_system(settings) as concurrency_system:
+        pool = AutoscaledPool(
+            concurrency_system=concurrency_system,
+            consumer=CONSUMER,
+            run_task_function=run,
+            is_task_ready_function=lambda: future(True),
+            is_finished_function=lambda: future(False),
+        )
+
+        pool_run_task = asyncio.create_task(pool.run(), name='pool run task')
+        try:
+            assert await poll_until_condition(lambda: pool.own_concurrency == 1)
+
+            pause_task = asyncio.create_task(pool.pause(), name='pause task')
+            await asyncio.sleep(0.1)
+            pool.resume()
+
+            first_release.set()
+            await asyncio.wait_for(pause_task, timeout=5)
+
+            assert await poll_until_condition(lambda: run.await_count == 2)
+            assert pool.own_concurrency == 1
+        finally:
+            first_release.set()
+            later_release.set()
+            pool_run_task.cancel()
+            await asyncio.gather(pool_run_task, return_exceptions=True)
+
+
+@pytest.mark.parametrize(
+    'is_ready',
+    [
+        pytest.param(True, id='task ready'),
+        pytest.param(False, id='no task ready'),
+    ],
+)
+async def test_pause_during_readiness_check(
+    make_concurrency_system: type[ConcurrencySystem], *, is_ready: bool
+) -> None:
+    """A pool paused while `is_task_ready_function` was awaited neither starts a task nor asks if it's finished."""
+    checking = asyncio.Event()
+    answer = asyncio.Event()
+
+    async def answer_after_pause() -> bool:
+        checking.set()
+        await answer.wait()
+        return is_ready
+
+    run = AsyncMock()
+    is_finished = AsyncMock(return_value=False)
+
+    async with make_concurrency_system() as concurrency_system:
+        pool = AutoscaledPool(
+            concurrency_system=concurrency_system,
+            consumer=CONSUMER,
+            run_task_function=run,
+            is_task_ready_function=AsyncMock(side_effect=answer_after_pause),
+            is_finished_function=is_finished,
+        )
+
+        pool_run_task = asyncio.create_task(pool.run(), name='pool run task')
+        try:
+            await asyncio.wait_for(checking.wait(), timeout=5)
+            await pool.pause()
+
+            answer.set()
+            await asyncio.sleep(0.1)
+            run.assert_not_awaited()
+            is_finished.assert_not_awaited()
+        finally:
+            pool_run_task.cancel()
+            await asyncio.gather(pool_run_task, return_exceptions=True)
+
+
+async def test_pause_during_finish_check(make_concurrency_system: type[ConcurrencySystem]) -> None:
+    """A pool paused while `is_finished_function` was awaited finishes only once it's resumed."""
+    checking = asyncio.Event()
+    answer = asyncio.Event()
+
+    async def answer_after_pause() -> bool:
+        checking.set()
+        await answer.wait()
+        return True
+
+    is_finished = AsyncMock(side_effect=answer_after_pause)
+
+    async with make_concurrency_system() as concurrency_system:
+        pool = AutoscaledPool(
+            concurrency_system=concurrency_system,
+            consumer=CONSUMER,
+            run_task_function=AsyncMock(),
+            is_task_ready_function=lambda: future(False),
+            is_finished_function=is_finished,
+        )
+
+        pool_run_task = asyncio.create_task(pool.run(), name='pool run task')
+        try:
+            await asyncio.wait_for(checking.wait(), timeout=5)
+            await pool.pause()
+
+            answer.set()
+            await asyncio.sleep(0.1)
+            assert not pool_run_task.done()
+
+            pool.resume()
+            await asyncio.wait_for(pool_run_task, timeout=5)
+        finally:
+            pool_run_task.cancel()
+            await asyncio.gather(pool_run_task, return_exceptions=True)
+
+    assert is_finished.await_count == 2
+
+
 @pytest.mark.run_alone
 async def test_runs_concurrently(make_concurrency_system: type[ConcurrencySystem]) -> None:
     """The pool fills the whole budget the system allows, rather than running the tasks one after another."""