diff --git a/loq.toml b/loq.toml index d06bc43..d0da5c4 100644 --- a/loq.toml +++ b/loq.toml @@ -10,7 +10,7 @@ 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" @@ -18,7 +18,7 @@ max_lines = 945 [[rules]] path = "src/docket/execution.py" -max_lines = 1335 +max_lines = 1336 [[rules]] path = "src/docket/docket.py" @@ -26,7 +26,7 @@ max_lines = 1098 [[rules]] path = "src/docket/_redis.py" -max_lines = 940 +max_lines = 969 [[rules]] path = "src/docket/dependencies/_concurrency.py" diff --git a/src/docket/_cancellation.py b/src/docket/_cancellation.py index f3f95eb..159ac40 100644 --- a/src/docket/_cancellation.py +++ b/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 diff --git a/src/docket/_execution_progress.py b/src/docket/_execution_progress.py index 13cae8d..7ad8ac4 100644 --- a/src/docket/_execution_progress.py +++ b/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 diff --git a/src/docket/_redis.py b/src/docket/_redis.py index aacbc05..18d2a50 100644 --- a/src/docket/_redis.py +++ b/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: diff --git a/src/docket/_redis_sentinel.py b/src/docket/_redis_sentinel.py index 8fd1b2c..854492b 100644 --- a/src/docket/_redis_sentinel.py +++ b/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, diff --git a/src/docket/execution.py b/src/docket/execution.py index 9ddec75..79c0aba 100644 --- a/src/docket/execution.py +++ b/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 diff --git a/src/docket/worker.py b/src/docket/worker.py index 567f714..cc0c7b3 100644 --- a/src/docket/worker.py +++ b/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 diff --git a/tests/_container.py b/tests/_container.py index e2b3b4a..3493352 100644 --- a/tests/_container.py +++ b/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) diff --git a/tests/test_automatic_perpetual_recovery.py b/tests/test_automatic_perpetual_recovery.py index 9c315d5..4f18f0a 100644 --- a/tests/test_automatic_perpetual_recovery.py +++ b/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") diff --git a/tests/test_execution_state.py b/tests/test_execution_state.py index d09180b..a5e4bd4 100644 --- a/tests/test_execution_state.py +++ b/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 diff --git a/tests/test_redelivery.py b/tests/test_redelivery.py index 2aa50e0..7fdee2e 100644 --- a/tests/test_redelivery.py +++ b/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), )