Fix shutdown regression: notify idle workers on stop()

Workers parked on job_available.wait() were not woken by stop(),
causing them to sleep out the full poll_idle_interval_s before
noticing _stop. With the default 1.0s interval, stop() took ~0.8s
instead of ~0.007s.

Notify all waiters on the condition in stop() so idle workers exit
immediately. Add tests for fast job pickup via notification and
fast shutdown with idle workers.
This commit is contained in:
Chris McDonough 2026-06-01 09:21:24 -04:00
parent 4fa8a012b4
commit 720ff23357
2 changed files with 45 additions and 0 deletions

View file

@ -92,6 +92,10 @@ class WorkerPool:
async def stop(self) -> None:
self._stop.set()
# Wake workers parked on job_available.wait() so they notice _stop
# immediately instead of sleeping out the full poll_idle interval.
async with self._jobs.job_available:
self._jobs.job_available.notify_all()
tasks = list(self._workers)
if self._reaper is not None:
tasks.append(self._reaper)

View file

@ -58,6 +58,47 @@ def _pool(client, jobs, sync, **kwargs) -> WorkerPool:
)
# --- event-driven wakeup ---
@pytest.mark.asyncio
async def test_idle_worker_picks_up_job_quickly(client, jobs, sync):
"""An idle worker should wake up well under poll_idle_s when a job is
enqueued, thanks to the job_available condition notification."""
client.create_document_from_source.return_value = Document(
id="d", content="x", uri="u", metadata={"md5": "m", "source_revision": "r"}
)
pool = _pool(client, jobs, sync, worker_count=1, poll_idle_interval_s=5.0)
await pool.start()
try:
await jobs.enqueue("src", "u", JobOp.UPSERT)
for _ in range(50):
listed = await jobs.list_jobs(status=JobStatus.SUCCEEDED)
if listed:
break
await asyncio.sleep(0.05)
assert len(listed) == 1, "job was not picked up within 2.5s"
finally:
await pool.stop()
@pytest.mark.asyncio
async def test_stop_completes_with_idle_workers(client, jobs, sync):
"""stop() must notify workers parked on job_available.wait() so they
exit promptly. Without the notify, workers sleep for the full
poll_idle_interval_s before noticing _stop."""
pool = _pool(client, jobs, sync, worker_count=2, poll_idle_interval_s=10.0)
await pool.start()
await asyncio.sleep(0.1)
try:
await asyncio.wait_for(pool.stop(), timeout=2.0)
except TimeoutError:
pytest.fail(
"stop() did not complete within 2s — idle workers were not woken"
)
assert pool.live_workers == 0
# --- drain_once: covers _process logic deterministically ---