Fix worker __aexit__ hanging on early cancellation (#263)

If a CancelledError arrived between _worker_done.clear() and entering
the
try block, the finally block never ran, leaving _worker_done cleared
forever.
This caused __aexit__ to hang indefinitely waiting on
_worker_done.wait().

The fix moves _cancellation_listener_task from __aenter__/__aexit__ to
_worker_loop, matching the pattern used by scheduler_task and
lease_renewal_task.
All async setup now happens inside the try block, ensuring the finally
block
always runs and _worker_done.set() is called.

Also names all asyncio tasks for easier debugging (docket.worker.*,
docket.strikelist.*).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: chrisguidry

src/docket/strikelist.py

@@ -227,7 +227,9 @@ async def connect(self) -> None:
         self._connection_pool = await connection_pool_from_url(self.url)
 
         self._strikes_loaded = asyncio.Event()
-        self._monitor_task = asyncio.create_task(self._monitor_strikes())
+        self._monitor_task = asyncio.create_task(
+            self._monitor_strikes(), name="docket.strikelist.monitor"
+        )
 
     async def close(self) -> None:
         """Close the Redis connection and stop monitoring.

src/docket/worker.py

@@ -205,16 +205,15 @@ def _maybe_suppress_instrumentation(self) -> Generator[None, None, None]:
             yield
 
     async def __aenter__(self) -> Self:
-        self._heartbeat_task = asyncio.create_task(self._heartbeat())
+        self._heartbeat_task = asyncio.create_task(
+            self._heartbeat(), name="docket.worker.heartbeat"
+        )
         self._execution_counts: dict[str, int] = {}
         # Track concurrency slots for active tasks so we can refresh them during
         # lease renewal. Maps execution.key → concurrency_key
         self._concurrency_slots: dict[str, str] = {}
         # Track running tasks for cancellation lookup
         self._tasks_by_key: dict[TaskKey, asyncio.Task[None]] = {}
-        self._cancellation_listener_task = asyncio.create_task(
-            self._cancellation_listener()
-        )
         # Events for coordinating worker loop shutdown
         self._worker_stopping = asyncio.Event()
         self._worker_done = asyncio.Event()
@@ -233,11 +232,6 @@ async def __aexit__(
         self._worker_stopping.set()
         await self._worker_done.wait()
 
-        self._cancellation_listener_task.cancel()
-        with suppress(asyncio.CancelledError):
-            await self._cancellation_listener_task
-        del self._cancellation_listener_task
-
         self._heartbeat_task.cancel()
         with suppress(asyncio.CancelledError):
             await self._heartbeat_task
@@ -346,10 +340,12 @@ def handle_shutdown(sig_name: str) -> None:  # pragma: no cover
 
                     try:
                         if until_finished:
-                            run_task = asyncio.create_task(worker.run_until_finished())
+                            run_task = asyncio.create_task(
+                                worker.run_until_finished(), name="docket.worker.run"
+                            )
                         else:
                             run_task = asyncio.create_task(
-                                worker.run_forever()
+                                worker.run_forever(), name="docket.worker.run"
                             )  # pragma: no cover
                         await run_task
                     except asyncio.CancelledError:  # pragma: no cover
@@ -419,21 +415,16 @@ async def _run(self, forever: bool = False) -> None:
     async def _worker_loop(self, redis: Redis, forever: bool = False):
         self._worker_stopping.clear()
         self._worker_done.clear()
+        self._cancellation_ready.clear()  # Reset for reconnection scenarios
 
-        await self._cancellation_ready.wait()
-
-        if self.schedule_automatic_tasks:
-            await self._schedule_all_automatic_perpetual_tasks()
-
+        # Initialize task variables before try block so finally can check them.
+        # This ensures _worker_done.set() is always called if _worker_done.clear() was.
+        cancellation_listener_task: asyncio.Task[None] | None = None
+        scheduler_task: asyncio.Task[None] | None = None
+        lease_renewal_task: asyncio.Task[None] | None = None
         active_tasks: dict[asyncio.Task[None], RedisMessageID] = {}
-
-        scheduler_task = asyncio.create_task(self._scheduler_loop(redis))
-        lease_renewal_task = asyncio.create_task(
-            self._renew_leases(redis, active_tasks)
-        )
         task_executions: dict[asyncio.Task[None], Execution] = {}
         available_slots = self.concurrency
-
         log_context = self._log_context()
 
         async def check_for_work() -> bool:
@@ -547,10 +538,27 @@ async def ack_message(redis: Redis, message_id: RedisMessageID) -> None:
                 )
                 await pipeline.execute()
 
-        has_work: bool = True
-        stopping = self._worker_stopping.is_set
-
         try:
+            # Start cancellation listener and wait for it to be ready
+            cancellation_listener_task = asyncio.create_task(
+                self._cancellation_listener(),
+                name="docket.worker.cancellation_listener",
+            )
+            await self._cancellation_ready.wait()
+
+            if self.schedule_automatic_tasks:
+                await self._schedule_all_automatic_perpetual_tasks()
+
+            scheduler_task = asyncio.create_task(
+                self._scheduler_loop(redis), name="docket.worker.scheduler"
+            )
+            lease_renewal_task = asyncio.create_task(
+                self._renew_leases(redis, active_tasks),
+                name="docket.worker.lease_renewal",
+            )
+
+            has_work: bool = True
+            stopping = self._worker_stopping.is_set
             while (forever or has_work or active_tasks) and not stopping():
                 await process_completed_tasks()
 
@@ -588,9 +596,21 @@ async def ack_message(redis: Redis, message_id: RedisMessageID) -> None:
                 await asyncio.gather(*active_tasks, return_exceptions=True)
                 await process_completed_tasks()
 
+            # Signal internal tasks to stop
             self._worker_stopping.set()
-            await scheduler_task
-            await lease_renewal_task
+
+            # These check _worker_stopping and exit cleanly
+            if scheduler_task is not None:
+                await scheduler_task
+            if lease_renewal_task is not None:
+                await lease_renewal_task
+
+            # Cancellation listener has while True loop, needs explicit cancellation
+            if cancellation_listener_task is not None:  # pragma: no branch
+                cancellation_listener_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await cancellation_listener_task
+
             self._worker_done.set()
 
     async def _scheduler_loop(self, redis: Redis) -> None:
@@ -1001,7 +1021,9 @@ async def _run_function_with_timeout(
                 },
             ),
         )
-        task = asyncio.create_task(task_coro)
+        task = asyncio.create_task(
+            task_coro, name=f"docket.worker.task:{execution.key}"
+        )
         try:
             while not task.done():  # pragma: no branch
                 remaining = timeout.remaining().total_seconds()

tests/test_worker.py

@@ -1,9 +1,9 @@
 import asyncio
 import logging
-import sys
-from contextlib import asynccontextmanager, suppress
+from contextlib import asynccontextmanager
 from datetime import datetime, timedelta, timezone
 from typing import AsyncGenerator, Callable
+
 from unittest.mock import AsyncMock, patch
 from uuid import uuid4
 
@@ -18,30 +18,13 @@
     Docket,
     Perpetual,
     Worker,
-    testing,
 )
 from docket.dependencies import Timeout
 from docket.execution import Execution
 from docket.tasks import standard_tasks
 from docket.worker import ms
 from tests._key_leak_checker import KeyCountChecker
 
-if sys.version_info >= (3, 11):  # pragma: no cover
-    from asyncio import timeout as async_timeout
-else:  # pragma: no cover
-
-    @asynccontextmanager
-    async def async_timeout(delay: float):
-        """Compatibility shim for asyncio.timeout on Python 3.10."""
-        task = asyncio.current_task()
-        loop = asyncio.get_running_loop()
-        deadline = loop.time() + delay
-        handle = loop.call_at(deadline, task.cancel)  # type: ignore[union-attr]
-        try:
-            yield
-        finally:
-            handle.cancel()
-
 
 async def test_worker_acknowledges_messages(
     docket: Docket, worker: Worker, the_task: AsyncMock
@@ -1135,74 +1118,3 @@ async def mock_xreadgroup(  # pyright: ignore[reportUnknownParameterType]
         assert task_executed
         # Should have called xreadgroup at least twice (once NOGROUP, then success)
         assert call_count >= 2
-
-
-async def test_run_forever_cancels_promptly_with_future_tasks(
-    docket: Docket, the_task: AsyncMock, now: Callable[[], datetime]
-):
-    """run_forever() should cancel promptly even with future-scheduled tasks.
-
-    Issue #260: Perpetual tasks block worker shutdown.
-    """
-    execution = await docket.add(the_task, now() + timedelta(seconds=15))()
-
-    async with Worker(
-        docket,
-        minimum_check_interval=timedelta(milliseconds=5),
-        scheduling_resolution=timedelta(milliseconds=5),
-    ) as worker:
-        worker_task = asyncio.create_task(worker.run_forever())
-        await asyncio.sleep(0.05)
-        worker_task.cancel()
-        with suppress(asyncio.CancelledError):  # pragma: no branch
-            async with async_timeout(1.0):  # pragma: no branch
-                await worker_task
-
-    the_task.assert_not_called()
-    await testing.assert_task_scheduled(docket, the_task, key=execution.key)
-
-
-async def test_run_until_finished_exits_promptly_with_future_tasks(
-    docket: Docket, the_task: AsyncMock, now: Callable[[], datetime]
-):
-    """run_until_finished() should exit promptly when only future tasks exist.
-
-    Issue #260: Perpetual tasks block worker shutdown.
-    """
-    execution = await docket.add(the_task, now() + timedelta(seconds=15))()
-
-    async with Worker(
-        docket,
-        minimum_check_interval=timedelta(milliseconds=5),
-        scheduling_resolution=timedelta(milliseconds=5),
-    ) as worker:
-        async with async_timeout(1.0):
-            await worker.run_until_finished()
-
-    the_task.assert_not_called()
-    await testing.assert_task_scheduled(docket, the_task, key=execution.key)
-
-
-async def test_run_at_most_cancels_promptly_with_future_tasks(
-    docket: Docket, the_task: AsyncMock, now: Callable[[], datetime]
-):
-    """run_at_most() should cancel promptly even with future-scheduled tasks.
-
-    Issue #260: Perpetual tasks block worker shutdown.
-    """
-    execution = await docket.add(the_task, now() + timedelta(seconds=15))()
-
-    async with Worker(
-        docket,
-        minimum_check_interval=timedelta(milliseconds=5),
-        scheduling_resolution=timedelta(milliseconds=5),
-    ) as worker:
-        worker_task = asyncio.create_task(worker.run_at_most({execution.key: 1}))
-        await asyncio.sleep(0.05)
-        worker_task.cancel()
-        with suppress(asyncio.CancelledError):  # pragma: no branch
-            async with async_timeout(1.0):  # pragma: no branch
-                await worker_task
-
-    the_task.assert_not_called()
-    await testing.assert_task_scheduled(docket, the_task, key=execution.key)