Confirm subscriptions before ready; bound TCP connects

loq.toml

@@ -10,23 +10,23 @@ max_lines = 750
 # Source files that still need exceptions above 750
 [[rules]]
 path = "src/docket/worker.py"
-max_lines = 1388
+max_lines = 1394
 
 [[rules]]
 path = "src/docket/cli/__init__.py"
 max_lines = 945
 
 [[rules]]
 path = "src/docket/execution.py"
-max_lines = 1335
+max_lines = 1336
 
 [[rules]]
 path = "src/docket/docket.py"
 max_lines = 1098
 
 [[rules]]
 path = "src/docket/_redis.py"
-max_lines = 940
+max_lines = 969
 
 [[rules]]
 path = "src/docket/dependencies/_concurrency.py"

src/docket/_cancellation.py

@@ -1,54 +1,33 @@
-"""Cancellation utilities for distinguishing internal vs external cancellation.
-
-When cancelling asyncio tasks, we often need to distinguish between cancellations
-we initiated (e.g., for timeout or cleanup) vs external cancellations (e.g., from
-a shutdown signal). Python 3.9+ supports passing a message to task.cancel(msg=...)
-which we use as a sentinel to identify our own cancellations.
-
-Note on Python version differences:
-- Python 3.11+: The cancel message propagates to the awaiter's CancelledError
-- Python 3.10: The message is stored but does NOT propagate to the awaiter
-
-The cancel_task() helper handles this by tracking that we initiated the cancel.
+"""Cancellation utilities for cleaning up background tasks.
+
+When we cancel one of our own background tasks (a heartbeat, a monitor, a
+renewal loop) and await it, the task's CancelledError surfaces in the awaiter
+and must not propagate: we asked for it.  A CancelledError should only
+propagate when it was delivered to the *caller* -- an external cancellation
+that arrived while we were waiting.
+
+Telling the two apart by the cancel message (task.cancel(msg=...)) is
+unreliable: CPython drops the message when the task's awaited future completes
+in the same event-loop tick as the cancellation (python/cpython#91048), and on
+Python 3.10 the message never propagates at all.  The awaited task's final
+state is reliable on every supported version, so cancel_task() checks that
+instead.
 """
 
 import asyncio
-import sys
 from typing import Any
 
 # Sentinel message for internal cancellation during cleanup
 CANCEL_MSG_CLEANUP = "docket:cleanup"
 
 
-def is_our_cancellation(exc: asyncio.CancelledError, expected_msg: str) -> bool:
-    """Check if we initiated this cancellation (vs. someone cancelling us).
-
-    When we cancel a task with task.cancel(msg), the CancelledError will have
-    args[0] set to that message. External cancellations (from TaskGroup, signals,
-    etc.) typically have no message or a different message.
-
-    Note: In Python 3.10, the message does NOT propagate to the awaiter, so this
-    function will return False even for our own cancellations. Use cancel_task()
-    instead for reliable cross-version behavior.
-
-    Args:
-        exc: The CancelledError to check
-        expected_msg: The sentinel message we used when cancelling
-
-    Returns:
-        True if the cancellation was initiated by us with the expected message
-    """
-    return bool(exc.args and exc.args[0] == expected_msg)
-
-
 async def cancel_task(task: "asyncio.Task[Any]", reason: str) -> None:
-    """Cancel a task and await its completion, suppressing our own cancellation.
-
-    This handles Python 3.10/3.11+ compatibility: in Python 3.10, the cancel
-    message doesn't propagate to the awaiter, so we can't rely on checking
-    the message. Instead, we track that we initiated the cancellation.
+    """Cancel a task and await its completion, suppressing its cancellation.
 
-    In Python 3.11+, we verify the message matches as an extra check.
+    If the awaited task ended cancelled, the CancelledError came from it, and
+    we swallow it -- we initiated that cancellation.  If the task did not end
+    cancelled, the CancelledError was delivered to the caller instead, so it
+    propagates.
 
     Args:
         task: The task to cancel
@@ -57,12 +36,6 @@ async def cancel_task(task: "asyncio.Task[Any]", reason: str) -> None:
     task.cancel(reason)
     try:
         await task
-    except asyncio.CancelledError as e:
-        if is_our_cancellation(e, reason):
-            return  # pragma: no cover - only on 3.11+ when message propagates
-        # Python 3.10: message doesn't propagate, but we just called cancel()
-        # so this CancelledError is probably from our cancel() call
-        if sys.version_info < (3, 11):  # pragma: no cover
-            return
-        # External cancellation - propagate it
-        raise  # pragma: no cover - race condition between cancel() and await
+    except asyncio.CancelledError:
+        if not task.cancelled():  # pragma: no cover - caller cancelled mid-await
+            raise

src/docket/_execution_progress.py

@@ -14,7 +14,7 @@
 )
 
 from ._lua import Arg, Args, Key, redis_script
-from ._redis import RedisClient
+from ._redis import RedisClient, confirm_subscriptions
 
 from ._telemetry import suppress_instrumentation
 from typing_extensions import Self
@@ -271,6 +271,7 @@ async def subscribe(
         channel = self.docket.key(f"progress:{self.key}")
         async with self.docket._pubsub() as pubsub:
             await pubsub.subscribe(channel)
+            await confirm_subscriptions(pubsub, 1)
             if ready is not None:
                 ready.set()
             async for message in pubsub.listen():  # pragma: no cover

src/docket/_redis.py

@@ -54,6 +54,14 @@
 # keepalive still detects genuinely dead peers.
 BLOCKING_READ_SOCKET_TIMEOUT: float | None = None
 
+# TCP connects are different: a connect that stalls (dropped SYNs on an
+# overloaded host) has no server-side work to wait for, so an unbounded hang
+# is never useful.  Without an explicit value, redis-py inherits the read
+# timeout above -- unbounded -- for connects too.  Bounding it lets redis-py's
+# retry layer surface a stalled connect as an error instead of hanging the
+# caller (a worker clearing its heartbeat during shutdown, for example).
+CONNECT_TIMEOUT: float = 10.0
+
 
 # ---------------------------------------------------------------------------
 # Input type aliases (mirror redis-py's type domain)
@@ -229,6 +237,22 @@ def listen(self) -> AsyncIterator[dict[str, Any]]: ...
     async def aclose(self) -> None: ...
 
 
+async def confirm_subscriptions(pubsub: PubSubClient, count: int) -> None:
+    """Block until the server has confirmed ``count`` subscription commands.
+
+    ``PubSub.subscribe()`` only sends the SUBSCRIBE command; the server's
+    confirmation arrives later as the first pub/sub message.  A publisher on
+    another connection can therefore beat the subscription unless the caller
+    waits for the confirmation before signaling readiness.
+
+    Subscription confirmations are delivered on a fresh connection before any
+    published message can arrive, so consuming the first ``count`` messages
+    consumes exactly the confirmations.
+    """
+    for _ in range(count):
+        await pubsub.get_message(ignore_subscribe_messages=False, timeout=None)
+
+
 # ---------------------------------------------------------------------------
 # RedisClient: the public surface advertised by docket.redis()
 # ---------------------------------------------------------------------------
@@ -791,7 +815,9 @@ async def _create_cluster_client(self) -> RedisCluster:  # pragma: no cover
             An initialized RedisCluster client ready for use
         """
         client: RedisCluster = RedisCluster.from_url(
-            self._normalized_url(), socket_timeout=BLOCKING_READ_SOCKET_TIMEOUT
+            self._normalized_url(),
+            socket_timeout=BLOCKING_READ_SOCKET_TIMEOUT,
+            socket_connect_timeout=CONNECT_TIMEOUT,
         )
         await client.initialize()
         return client
@@ -821,6 +847,7 @@ def _create_node_pool(self) -> ConnectionPool:  # pragma: no cover
             else Connection,
             decode_responses=False,
             socket_timeout=BLOCKING_READ_SOCKET_TIMEOUT,
+            socket_connect_timeout=CONNECT_TIMEOUT,
         )
 
     async def _connection_pool_from_url(
@@ -845,11 +872,13 @@ async def _connection_pool_from_url(
                 self.url,
                 decode_responses=decode_responses,
                 socket_timeout=BLOCKING_READ_SOCKET_TIMEOUT,
+                socket_connect_timeout=CONNECT_TIMEOUT,
             )
         return ConnectionPool.from_url(  # pyright: ignore[reportUnknownMemberType]
             self.url,
             decode_responses=decode_responses,
             socket_timeout=BLOCKING_READ_SOCKET_TIMEOUT,
+            socket_connect_timeout=CONNECT_TIMEOUT,
         )
 
     async def _get_or_create_memory_client(self) -> MemoryRedisClient:

src/docket/_redis_sentinel.py

@@ -256,19 +256,26 @@ async def aclose(self) -> None:
 
 
 def sentinel_connection_pool(
-    url: str, *, decode_responses: bool, socket_timeout: float | None
+    url: str,
+    *,
+    decode_responses: bool,
+    socket_timeout: float | None,
+    socket_connect_timeout: float | None,
 ) -> ConnectionPool:
     """Create a connection pool that resolves the master through Sentinel.
 
     The pool asks the listed Sentinel daemons for the current master and follows
     failover automatically, so the rest of the standalone code path (client,
     pub/sub, publish, result storage) works unchanged.  ``socket_timeout`` is
-    docket's blocking-read timeout, applied to the data-node connections.
+    docket's blocking-read timeout and ``socket_connect_timeout`` is its
+    bounded TCP connect timeout, both applied to the data-node connections.
 
     Args:
         url: The redis+sentinel:// or rediss+sentinel:// URL.
         decode_responses: If True, decode Redis responses from bytes to strings.
         socket_timeout: The read timeout for data-node connections.
+        socket_connect_timeout: The TCP connect timeout for data-node
+            connections.
 
     Returns:
         A ConnectionPool ready for use with Redis clients.
@@ -281,6 +288,7 @@ def sentinel_connection_pool(
         "db": config.db,
         "decode_responses": decode_responses,
         "socket_timeout": socket_timeout,
+        "socket_connect_timeout": socket_connect_timeout,
         "socket_keepalive": True,
         "socket_keepalive_options": SENTINEL_SOCKET_KEEPALIVE_OPTIONS,
         **config.connection_kwargs,

src/docket/execution.py

@@ -36,7 +36,7 @@
 
 from ._execution_progress import ExecutionProgress, ProgressEvent, StateEvent
 from ._lua import Arg, Args, Key, redis_script
-from ._redis import RedisClient, is_cluster_client
+from ._redis import RedisClient, confirm_subscriptions, is_cluster_client
 from .annotations import Logged
 from .instrumentation import CACHE_SIZE, message_getter, message_setter
 
@@ -1296,6 +1296,7 @@ async def subscribe(
         progress_channel = self.docket.key(f"progress:{self.key}")
         async with self.docket._pubsub() as pubsub:
             await pubsub.subscribe(state_channel, progress_channel)
+            await confirm_subscriptions(pubsub, 2)
             if ready is not None:
                 ready.set()
             async for message in pubsub.listen():  # pragma: no cover

src/docket/worker.py

@@ -1346,13 +1346,19 @@ async def _cancellation_listener(self) -> None:
             try:
                 async with self.docket._pubsub() as pubsub:
                     await pubsub.psubscribe(cancel_pattern)
-                    session.cancellation_ready.set()
                     while not session.stopping.is_set():
                         message = await pubsub.get_message(
-                            ignore_subscribe_messages=True, timeout=0.1
+                            ignore_subscribe_messages=False, timeout=0.1
                         )
-                        if message is not None and message["type"] == "pmessage":
+                        if message is None:
+                            continue
+                        if message["type"] == "pmessage":
                             await self._handle_cancellation(message)
+                        else:
+                            # The PSUBSCRIBE confirmation: the server now has
+                            # the subscription, so cancellations published
+                            # from here on will be delivered.
+                            session.cancellation_ready.set()
             except ConnectionError:
                 if session.stopping.is_set():
                     return  # pragma: no cover

tests/_container.py

@@ -55,8 +55,11 @@ def __init__(self, worker_id: str) -> None:
 
 @contextmanager
 def sync_redis(url: str) -> Generator[Redis, None, None]:
+    # redis-py 8 defaults the sync client's socket_timeout to 5s, which a
+    # FLUSHALL on a busy CI container can exceed.  Keep reads bounded but
+    # generous so fixture setup fails only when Redis is genuinely stuck.
     pool: ConnectionPool | None = None
-    redis = Redis.from_url(url)  # type: ignore
+    redis = Redis.from_url(url, socket_timeout=30)  # type: ignore
     try:
         with redis:
             pool = redis.connection_pool
@@ -82,7 +85,7 @@ def wait_for_redis(port: int) -> None:
             with administrative_redis(port) as r:
                 if r.ping():  # type: ignore
                     return
-        except redis.exceptions.ConnectionError:
+        except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError):
             time.sleep(0.1)
 
 

tests/test_automatic_perpetual_recovery.py

@@ -62,8 +62,13 @@ async def next_run_is_parked() -> bool:
     await wait_until(next_run_is_parked, description="next run parked")
 
     # Sever the chain the same way the fallback drop does: no scheduled entry,
-    # nothing in flight.
-    await docket.cancel("perpetual_task")
+    # nothing in flight.  Work directly on the keys rather than through
+    # docket.cancel(), whose pub/sub signal can cancel the first execution's
+    # still-running bookkeeping and race the state this test asserts on.
+    async with docket.redis() as redis:
+        await redis.zrem(docket.queue_key, "perpetual_task")
+        await redis.delete(docket.parked_task_key("perpetual_task"))
+        await redis.hdel(runs_key, "known", "stream_id")
 
     await wait_until(lambda: calls == 2, description="re-seeding re-ran the perpetual")
 

tests/test_execution_state.py

@@ -121,8 +121,10 @@ async def tracking_task(progress: Progress = Progress()):
             await progress.set_message(f"Step {i + 1}")
             await asyncio.sleep(0.01)
 
-    # Schedule task in the future
-    when = datetime.now(timezone.utc) + timedelta(milliseconds=50)
+    # Schedule the task far enough out that a slow CI runner can't cross the
+    # deadline between computing `when` here and the server evaluating it,
+    # which would enqueue the task immediately as QUEUED instead of SCHEDULED.
+    when = datetime.now(timezone.utc) + timedelta(seconds=2)
     execution = await docket.add(tracking_task, when=when)()
 
     # Should be SCHEDULED

tests/test_redelivery.py

@@ -141,18 +141,21 @@ async def test_multiple_workers_no_duplicate_execution(docket: Docket):
     async def slow_task(task_id: int, worker_name: str = ""):
         async with lock:
             executions.append((worker_name, task_id))
-        await asyncio.sleep(0.5)
+        await asyncio.sleep(1.5)
 
     # Schedule several tasks
     for i in range(6):
         await docket.add(slow_task, key=f"task-{i}")(task_id=i)
 
-    # Run multiple workers concurrently with short redelivery_timeout
+    # Tasks sleep longer than the redelivery_timeout, so without renewal a
+    # competing worker's XAUTOCLAIM would reclaim them.  Use a 1s timeout so
+    # renewal (every 250ms) tolerates the multi-hundred-millisecond stalls
+    # that CI runners hit; at 200ms a single late renewal caused a steal.
     workers = [
         Worker(
             docket,
             name=f"worker-{i}",
-            redelivery_timeout=timedelta(milliseconds=200),
+            redelivery_timeout=timedelta(seconds=1),
             minimum_check_interval=timedelta(milliseconds=10),
             scheduling_resolution=timedelta(milliseconds=10),
             concurrency=2,
@@ -295,15 +298,20 @@ async def test_worker_joining_doesnt_steal_renewed_lease(docket: Docket):
     async def slow_task(task_id: int):
         executions.append(("start", task_id))
         task_started.set()
-        await asyncio.sleep(0.6)  # Long task
+        # Sleep well past the redelivery_timeout so B's XAUTOCLAIM would
+        # reclaim the task if A's renewal ever let it go idle.
+        await asyncio.sleep(2.5)
         executions.append(("end", task_id))
 
     await docket.add(slow_task, key="task")(task_id=1)
 
+    # A 1s timeout gives renewal (every 250ms) enough margin for the
+    # multi-hundred-millisecond stalls that CI runners hit; at 200ms a
+    # single late renewal let worker B steal the task.
     worker_a = Worker(
         docket,
         name="worker-a",
-        redelivery_timeout=timedelta(milliseconds=200),
+        redelivery_timeout=timedelta(seconds=1),
         minimum_check_interval=timedelta(milliseconds=10),
         scheduling_resolution=timedelta(milliseconds=10),
     )
@@ -325,7 +333,7 @@ async def run_a():
     worker_b = Worker(
         docket,
         name="worker-b",
-        redelivery_timeout=timedelta(milliseconds=200),
+        redelivery_timeout=timedelta(seconds=1),
         minimum_check_interval=timedelta(milliseconds=10),
         scheduling_resolution=timedelta(milliseconds=10),
     )