fix(scheduler): drain owned tasks before shutdown

Author: Kuang-xianxin

taskiq/cli/scheduler/run.py

@@ -196,6 +196,8 @@ def __init__(
         self._update_schedules_task_future: asyncio.Task[Any] | None = None
 
     def _update_schedules_task_future_callback(self, task_: asyncio.Task[Any]) -> None:
+        if task_.cancelled():
+            return
         self.scheduled_tasks = task_.result()
 
         new_schedules_ids: set[ScheduleId] = set()
@@ -229,6 +231,18 @@ async def _update_scheduled_tasks(self) -> None:
                 self._update_schedules_task_future_callback,
             )
 
+    async def _cancel_pending_tasks(
+        self,
+        running_schedules: dict[ScheduleId, asyncio.Task[Any]],
+    ) -> None:
+        # Stop owned operations before callers close the broker and sources.
+        pending_tasks = list(running_schedules.values())
+        if self._update_schedules_task_future is not None:
+            pending_tasks.append(self._update_schedules_task_future)
+        for pending_task in pending_tasks:
+            pending_task.cancel()
+        await asyncio.gather(*pending_tasks, return_exceptions=True)
+
     def _mark_cron_tasks_as_already_run(self) -> None:
         current_minute = datetime.now(tz=timezone.utc).replace(second=0, microsecond=0)
         for _, task_list in self.scheduled_tasks:
@@ -314,42 +328,48 @@ async def run(
 
         await _sleep_until_next_second()
 
-        while True:
-            now = datetime.now(tz=timezone.utc)
-            next_run = (now + loop_interval).replace(microsecond=0)
-
-            if now - self.scheduled_tasks_updated_at >= update_interval:
-                await self._update_scheduled_tasks()
-                self.scheduled_tasks_updated_at = now
-
-            for source, task_list in self.scheduled_tasks:
-                for task in task_list:
-                    is_ready_to_send: bool = self._is_schedule_ready_to_send(
-                        task=task,
-                        now=now,
-                    )
-
-                    if is_ready_to_send and task.schedule_id not in running_schedules:
-                        send_task = self._event_loop.create_task(
-                            send(self.scheduler, source, task),
-                            # We need to set the name of the task
-                            # to be able to discard its reference
-                            # after it is done.
-                            name=f"schedule_{task.schedule_id}",
-                        )
-                        running_schedules[task.schedule_id] = send_task
-                        send_task.add_done_callback(
-                            lambda task_future: running_schedules.pop(
-                                task_future.get_name().removeprefix("schedule_"),
-                            ),
+        try:
+            while True:
+                now = datetime.now(tz=timezone.utc)
+                next_run = (now + loop_interval).replace(microsecond=0)
+
+                if now - self.scheduled_tasks_updated_at >= update_interval:
+                    await self._update_scheduled_tasks()
+                    self.scheduled_tasks_updated_at = now
+
+                for source, task_list in self.scheduled_tasks:
+                    for task in task_list:
+                        is_ready_to_send: bool = self._is_schedule_ready_to_send(
+                            task=task,
+                            now=now,
                         )
 
-            delay = next_run - datetime.now(tz=timezone.utc)
-            logger.debug(
-                "Sleeping for %.3f seconds before getting schedules.",
-                delay.total_seconds(),
-            )
-            await asyncio.sleep(delay.total_seconds())
+                        if (
+                            is_ready_to_send
+                            and task.schedule_id not in running_schedules
+                        ):
+                            send_task = self._event_loop.create_task(
+                                send(self.scheduler, source, task),
+                                # We need to set the name of the task
+                                # to be able to discard its reference
+                                # after it is done.
+                                name=f"schedule_{task.schedule_id}",
+                            )
+                            running_schedules[task.schedule_id] = send_task
+                            send_task.add_done_callback(
+                                lambda task_future: running_schedules.pop(
+                                    task_future.get_name().removeprefix("schedule_"),
+                                ),
+                            )
+
+                delay = next_run - datetime.now(tz=timezone.utc)
+                logger.debug(
+                    "Sleeping for %.3f seconds before getting schedules.",
+                    delay.total_seconds(),
+                )
+                await asyncio.sleep(delay.total_seconds())
+        finally:
+            await self._cancel_pending_tasks(running_schedules)
 
 
 async def run_scheduler(args: SchedulerArgs) -> None:

tests/cli/scheduler/test_shutdown.py

@@ -0,0 +1,140 @@
+import asyncio
+from contextlib import suppress
+from datetime import timedelta
+from typing import Any
+from unittest.mock import AsyncMock
+
+import pytest
+
+from taskiq import BrokerMessage, ScheduleSource, TaskiqScheduler
+from taskiq.api import run_scheduler_task
+from taskiq.cli.scheduler.args import SchedulerArgs
+from taskiq.cli.scheduler.run import SchedulerLoop, run_scheduler
+from taskiq.scheduler.scheduled_task import ScheduledTask
+from tests.utils import AsyncQueueBroker
+
+
+class PendingOperation:
+    def __init__(self, cleanup_allowed: asyncio.Event) -> None:
+        self.started = asyncio.Event()
+        self.cleanup_started = asyncio.Event()
+        self.finished = asyncio.Event()
+        self.cleanup_allowed = cleanup_allowed
+        self.task: asyncio.Task[Any] | None = None
+
+    async def run(self) -> None:
+        self.task = asyncio.current_task()
+        self.started.set()
+        try:
+            await asyncio.Event().wait()
+        finally:
+            self.cleanup_started.set()
+            await self.cleanup_allowed.wait()
+            self.finished.set()
+
+
+@pytest.mark.parametrize("entrypoint", ["loop", "api", "cli"])
+async def test_shutdown_drains_sends_and_updates(
+    entrypoint: str,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Owned operations must finish cleanup before scheduler resources close."""
+    cleanup_allowed = asyncio.Event()
+    sending = PendingOperation(cleanup_allowed)
+    updating = PendingOperation(cleanup_allowed)
+    shutdowns: list[str] = []
+
+    class BlockingBroker(AsyncQueueBroker):
+        async def kick(self, message: BrokerMessage) -> None:
+            await sending.run()
+
+        async def shutdown(self) -> None:
+            shutdowns.append("broker")
+            assert sending.finished.is_set()
+            assert updating.finished.is_set()
+            await super().shutdown()
+
+    class BlockingSource(ScheduleSource):
+        fetched = False
+
+        async def get_schedules(self) -> list[ScheduledTask]:
+            if self.fetched:
+                await updating.run()
+            self.fetched = True
+            return [
+                ScheduledTask(
+                    task_name="test",
+                    labels={},
+                    args=[],
+                    kwargs={},
+                    interval=1,
+                ),
+            ]
+
+        async def shutdown(self) -> None:
+            shutdowns.append("source")
+            assert updating.finished.is_set()
+
+    monkeypatch.setattr(
+        "taskiq.cli.scheduler.run._sleep_until_next_second",
+        AsyncMock(),
+    )
+    scheduler = TaskiqScheduler(BlockingBroker(), [BlockingSource()])
+    if entrypoint == "cli":
+        coroutine = run_scheduler(
+            SchedulerArgs(
+                scheduler=scheduler,
+                modules=[],
+                update_interval=0,
+                configure_logging=False,
+            ),
+        )
+    elif entrypoint == "api":
+        coroutine = run_scheduler_task(scheduler, interval=timedelta(0))
+    else:
+        coroutine = SchedulerLoop(scheduler).run(update_interval=timedelta(0))
+
+    scheduler_task = asyncio.create_task(coroutine)
+    cleanup_started = asyncio.gather(
+        sending.cleanup_started.wait(),
+        updating.cleanup_started.wait(),
+    )
+    try:
+        await asyncio.wait_for(
+            asyncio.gather(sending.started.wait(), updating.started.wait()),
+            5,
+        )
+        scheduler_task.cancel()
+        waitables: list[asyncio.Future[Any]] = [scheduler_task, cleanup_started]
+        done, _ = await asyncio.wait(
+            waitables,
+            timeout=5,
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+        assert cleanup_started in done
+        assert not scheduler_task.done()
+        assert not shutdowns
+
+        cleanup_allowed.set()
+        with suppress(asyncio.CancelledError):
+            await asyncio.wait_for(scheduler_task, 5)
+
+        assert sending.finished.is_set()
+        assert updating.finished.is_set()
+        assert scheduler_task.cancelled() == (entrypoint != "cli")
+        assert shutdowns == (["broker", "source"] if entrypoint == "cli" else [])
+    finally:
+        cleanup_allowed.set()
+        tasks = [
+            task
+            for task in (scheduler_task, sending.task, updating.task)
+            if task is not None
+        ]
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        cleanup_started.cancel()
+        await asyncio.gather(cleanup_started, return_exceptions=True)
+
+    assert "Exception in callback" not in caplog.text