Confirm subscriptions before ready; bound TCP connects

[chrisguidry]

CI failed on one or two jobs from time to time, and the same jobs passed on a re-run. I pulled the failure logs from the last 200 runs. They show three separate causes, and this change fixes all three. The re-run sequence for this PR then surfaced two more flakes, and this change fixes those too — seven in total.

The pub/sub acknowledgment race

The pub/sub tests timed out on every real Redis backend, in runs 29753363584, 29753344632, 28542779700, and 28188728646. The cause is a race in the library, not in the tests. PubSub.subscribe() only sends the SUBSCRIBE command; it does not wait for the reply from the server. We set the ready event immediately, so a publish on another connection could win the race against the subscription, and the subscriber lost those events.

The new confirm_subscriptions() helper consumes the confirmation messages from the server before it sets the ready event. Today's listen() loop already received those confirmation frames and discarded them, so the message flow does not change. The worker's cancellation listener gets the same guarantee inside its 0.1s poll loop, not with a blocking wait, so worker shutdown stays responsive.

Unbounded TCP connects

A CLI test hit its 30s timeout in run 28913775528. The stack shows an unbounded TCP connect inside Worker._remove_heartbeat during cleanup. Our pools set no socket_connect_timeout, so a stalled connect could hang forever. The new CONNECT_TIMEOUT bounds connects to 10s on all the pools. Blocking reads stay unbounded, as before.

The sync test client's 5s read timeout

Two tests got fixture errors in run 28071483236, with a sync redis.exceptions.TimeoutError. redis-py 8 gives the sync client a 5s read timeout by default, and a FLUSHALL on a busy CI container can exceed that. The sync_redis() test helper now uses socket_timeout=30, and wait_for_redis() also retries on TimeoutError.

The lost cancel message in cancel_task()

Attempt 4 of run 29757754279 failed test_snapshot_with_stats_flag_mixed_tasks with a bare CancelledError and a TTL leak. cancel_task() decided between "our cancellation" and "an external cancellation" by the message on the CancelledError. CPython drops that message when the awaited future completes in the same event-loop tick as the cancellation (python/cpython#91048). The worker's cleanup set session.stopping and then cancelled the heartbeat task, so the two raced, the message was lost, and the helper re-raised our own cleanup cancellation. That abort also skipped the drain of the active tasks, which leaked two progress:* keys. The helper now checks task.cancelled() on the awaited task, which is reliable on every supported Python version.

A 50ms margin in the lifecycle test

The same attempt failed test_full_lifecycle_integration, which scheduled a task 50ms in the future and then asserted that its state was SCHEDULED. On a slow runner, more than 50ms passed before the server evaluated the schedule, so the task enqueued immediately as QUEUED. The test now schedules 2 seconds out.

A 200ms lease in the steal tests

Attempt 1 of run 29760321411 failed test_worker_joining_doesnt_steal_renewed_lease on the Redis 8 Cluster job: the task ran twice. The test gave the workers a 200ms redelivery_timeout, and renewal runs at a quarter of that, so one renewal 150ms late lets a competing worker's XAUTOCLAIM steal the task. test_workers_with_same_redelivery_timeout already moved to a 1s timeout for the same reason; the two tests where a competing worker can steal now follow that precedent, with task sleeps that still exceed the timeout so the tests keep their meaning.

A cancel() that races the worker in the re-seed test

Attempt 2 of the same run failed test_running_worker_resows_severed_automatic_perpetual on Windows: the re-seed never re-ran the task within 5 seconds. The test severed the perpetual chain with docket.cancel(), whose pub/sub signal reached the worker while the first execution was still in its post-completion bookkeeping — the log shows "Cancelling running task" and no completion line. The test now severs the chain directly on the Redis keys, which is the exact state its comment describes ("no scheduled entry, nothing in flight") with no signal racing the live worker.

Verification

I verified the pub/sub fix with pytest-flakefinder: 80 runs on Redis 7.4, and 40 runs on Redis 6.2 at 0.2 CPU, all green. The core suite and the CLI suite both pass at 100% coverage. The pub/sub flake fired about weekly, so the real proof is a quiet couple of weeks in CI; I am re-running this PR's CI until it is green five times in a row.

Session context

Session summary: Chris asked me to find the tests that flake in CI and to fix them for good. I searched the last 200 GitHub Actions runs and pulled the failed-step logs from every failed run. Three flakes explained all the one-or-two-job failures; the broad multi-job failures were real bugs on development branches. The dominant flake was the pub/sub subscribe-acknowledgment race. I confirmed in redis-py's source that PubSub.execute_command only sends the command, and in burner-redis's source that it queues confirmation dicts at subscribe time, so one design works on all backends.

Major decisions:

• Chris asked me to validate that the acknowledgment wait cannot deadlock. The two subscribe() generators are safe: the wait is one round trip on a dedicated connection, a dead connection raises, and cancellation propagates. But a blocking wait in the worker's cancellation listener could hang worker shutdown, because the worker's TaskGroup waits for the listener to see session.stopping. So the listener keeps its 0.1s poll loop and folds the acknowledgment into it.
• Chris did not want large timeout raises, so all the test timeouts stay exactly as they were. The fixes remove the races instead.
• Chris set a goal of five flake-free CI runs in a row. Attempt 4 of the re-run sequence failed with two new flakes, which led to the cancel_task() fix and the lifecycle-test margin fix in the second commit. The count restarted after that commit.
• loq baselines for three files grew by the size of the additions, per the project's baseline workflow.

🤖 Generated with Claude Code

[read-the-docs-community]

Documentation build overview

📚 docket | 🛠️ Build #33671646 | 📁 Comparing 6124846 against latest (457830d)

  🔍 Preview build  

1 file changed

± api-reference/index.html

[codecov-commenter]

❌ 1 Tests Failed:

Tests completed	Failed	Passed	Skipped
958	1	957	7

View the top 2 failed test(s) by shortest run time

tests/test_execution_state.py::test_full_lifecycle_integration

Stack Traces | 0.206s run time

docket = <docket.docket.Docket object at 0x7fce0afe1550>
worker = <docket.worker.Worker object at 0x7fce1024b460>

    async def test_full_lifecycle_integration(docket: Docket, worker: Worker):
        """Test complete lifecycle: SCHEDULED -> QUEUED -> RUNNING -> COMPLETED."""
        states_observed: list[ExecutionState] = []
    
        async def tracking_task(progress: Progress = Progress()):
            await progress.set_total(3)
            for i in range(3):
                await progress.increment()
                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)
        execution = await docket.add(tracking_task, when=when)()
    
        # Should be SCHEDULED
        await execution.sync()
>       assert execution.state == ExecutionState.SCHEDULED
E       AssertionError: assert <ExecutionState.QUEUED: 'queued'> == <ExecutionState.SCHEDULED: 'scheduled'>
E        +  where <ExecutionState.QUEUED: 'queued'> = <docket.execution.Execution object at 0x7fce1024b8a0>.state
E        +  and   <ExecutionState.SCHEDULED: 'scheduled'> = ExecutionState.SCHEDULED

tests/test_execution_state.py:130: AssertionError

tests/cli/test_snapshot.py::test_snapshot_with_stats_flag_mixed_tasks

Stack Traces | 6.7s run time

docket = <docket.docket.Docket object at 0x7fb4453cbcb0>
key_leak_checker = <tests._key_leak_checker.KeyCountChecker object at 0x7fb4453c6c00>

    async def test_snapshot_with_stats_flag_mixed_tasks(
        docket: Docket, key_leak_checker: KeyCountChecker
    ):
        """Should show task count statistics when --stats flag is used"""
        heartbeat = timedelta(milliseconds=20)
        docket.heartbeat_interval = heartbeat
    
        # Add multiple tasks of different types
        future = datetime.now(timezone.utc) + timedelta(seconds=5)
        await docket.add(tasks.trace, when=future, key="trace-1")("hi!")
        await docket.add(tasks.trace, when=future, key="trace-2")("hello!")
    
        # This test cancels worker mid-execution, leaving incomplete tasks
        key_leak_checker.add_exemption(f"{docket.name}:runs:trace-1")
        key_leak_checker.add_exemption(f"{docket.name}:runs:trace-2")
    
        # Use tasks.sleep for CLI tests since the subprocess needs importable tasks
        for i in range(3):
            await docket.add(tasks.sleep, key=f"sleep-{i}")(4)
            key_leak_checker.add_exemption(f"{docket.name}:runs:sleep-{i}")
    
        async with Worker(docket, name="test-worker", concurrency=2) as worker:
            worker_running = asyncio.create_task(worker.run_until_finished())
    
            await _running_count_at_least(docket, 2)
    
            result = await run_cli(
                "snapshot",
                "--stats",
                "--url",
                docket.url,
                "--docket",
                docket.name,
            )
            assert result.exit_code == 0, result.output
    
            # Should show the normal summary
            assert "1 workers, 2/5 running" in result.output, result.output
    
            # Should show task statistics table with enhanced columns
            assert "Task Count Statistics by Function" in result.output
            assert "Function" in result.output
            assert "Total" in result.output
            assert "Running" in result.output
            assert "Queued" in result.output
            assert "Oldest Queued" in result.output
            assert "Latest Queued" in result.output
    
            # Should show the task counts
            assert "sleep" in result.output
            assert "trace" in result.output
    
            worker_running.cancel()
>           await worker_running

tests/cli/test_snapshot.py:273: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
src/docket/worker.py:498: in run_until_finished
    return await self._run(forever=False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/docket/worker.py:553: in _run
    return await self._worker_loop(redis, forever=forever)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/docket/worker.py:839: in _worker_loop
    await cancel_task(self._heartbeat_task, CANCEL_MSG_CLEANUP)
src/docket/_cancellation.py:59: in cancel_task
    await task
src/docket/worker.py:1333: in _heartbeat
    if await _wait_for_event(
src/docket/worker.py:114: in _wait_for_event
    await asyncio.wait_for(event.wait(), timeout=timeout)
....../usr/lib/python3.12/asyncio/tasks.py:520: in wait_for
    return await fut
           ^^^^^^^^^
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <asyncio.locks.Event object at 0x7fb44529b5c0 [set]>

    async def wait(self):
        """Block until the internal flag is true.
    
        If the internal flag is true on entry, return True
        immediately.  Otherwise, block until another coroutine calls
        set() to set the flag to true, then return True.
        """
        if self._value:
            return True
    
        fut = self._get_loop().create_future()
        self._waiters.append(fut)
        try:
>           await fut
E           asyncio.exceptions.CancelledError

....../usr/lib/python3.12/asyncio/locks.py:212: CancelledError

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}