Make ingester worker ids globally unique

This commit is contained in:
Yiorgis Gozadinos 2026-06-25 11:44:22 +03:00
parent 6bbf12e794
commit ba6b318ece
No known key found for this signature in database
3 changed files with 44 additions and 1 deletions

View file

@ -4,6 +4,7 @@
### Fixed
- A failed FTS index build is logged at `WARNING` instead of `DEBUG`, so silent full-text search degradation is visible.
- Ingester worker ids are now globally unique (`{pid}-{uuid}-{n}`) instead of `worker-{n}`, so the `claimed_by` guards on job completion/reschedule/release distinguish workers across processes sharing one Postgres queue.
### Changed

View file

@ -1,7 +1,9 @@
import asyncio
import logging
import os
import time
from typing import TYPE_CHECKING
from uuid import uuid4
from haiku.rag.config import CircuitBreakerConfig
from haiku.rag.ingester.exceptions import PermanentError, TransientError
@ -59,12 +61,20 @@ class WorkerPool:
self._metadata_providers: dict[str, MetadataProvider] = (
dict(metadata_providers) if metadata_providers else {}
)
# Globally-unique so claimed_by distinguishes this pool's workers from
# those of any other process sharing the queue; the claimed_by guards on
# mark_succeeded/reschedule/release rely on it. The uuid guarantees
# uniqueness; the pid just makes claimed_by readable in logs/dashboard.
self._instance = f"{os.getpid()}-{uuid4().hex[:8]}"
self._stop = asyncio.Event()
self._workers: list[asyncio.Task] = []
self._reaper: asyncio.Task | None = None
self._pending_releases: set[asyncio.Task] = set()
self._breakers: dict[str, CircuitBreaker] = {}
def _worker_id(self, i: int) -> str:
return f"{self._instance}-{i}"
@property
def live_workers(self) -> int:
"""Worker tasks that are still running. Equal to worker_count under
@ -106,7 +116,9 @@ class WorkerPool:
if reset:
logger.info("Boot-reaped %d stale claim(s) from previous process", reset)
for i in range(self._worker_count):
self._workers.append(asyncio.create_task(self._worker_loop(f"worker-{i}")))
self._workers.append(
asyncio.create_task(self._worker_loop(self._worker_id(i)))
)
self._reaper = asyncio.create_task(self._reaper_loop())
async def stop(self) -> None:

View file

@ -31,6 +31,36 @@ def _pool(client, jobs, sync, **kwargs) -> WorkerPool:
)
# --- worker identity ---
@pytest.mark.asyncio
async def test_worker_ids_are_unique_across_pools(client, jobs, sync):
"""Two pools built with default construction must not share worker ids;
otherwise a stale worker from one pool can satisfy the claimed_by guard of
a job re-claimed by another pool and clobber its result."""
pool_a = _pool(client, jobs, sync)
pool_b = _pool(client, jobs, sync)
worker_a = pool_a._worker_id(0)
worker_b = pool_b._worker_id(0)
assert worker_a != worker_b
job = await jobs.enqueue("src", "u", JobOp.UPSERT)
assert job is not None
claimed = await jobs.claim_next(worker_a)
assert claimed is not None
# Reaper resets the claim; pool B re-claims and finishes first.
await jobs.reap_stale(claim_timeout_seconds=0)
reclaimed = await jobs.claim_next(worker_b)
assert reclaimed is not None and reclaimed.id == job.id
assert await jobs.mark_succeeded(reclaimed.id, worker_b) is True
# Pool A's slow worker finishes later: with unique ids this is a no-op, so
# pool B's success is not clobbered.
assert await jobs.mark_succeeded(job.id, worker_a) is False
# --- event-driven wakeup ---