Coverage for src/qdrant_loader/core/worker/pool.py: 85%
168 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1from __future__ import annotations
3import asyncio
4import json
5import time
6from json import JSONDecodeError
7from typing import Any
9from qdrant_loader.core.worker.handlers import JobHandler, PermanentJobError
10from qdrant_loader.core.worker.queue import JobQueue
11from qdrant_loader.utils.logging import LoggingConfig
13logger = LoggingConfig.get_logger(__name__)
16class QueueWorkerPool:
17 """Run queue jobs with bounded concurrency and per-source serialization."""
19 DEFAULT_SOURCE_KEY = "__global__"
21 def __init__(
22 self,
23 queue: JobQueue,
24 handler: JobHandler,
25 worker_count: int = 4,
26 lease_seconds: int = 60,
27 max_attempts: int = 1,
28 retry_backoff_base_seconds: int = 0,
29 job_types: list[str] | None = None,
30 ) -> None:
31 if worker_count < 1:
32 raise ValueError("worker_count must be >= 1")
33 if lease_seconds < 1:
34 raise ValueError("lease_seconds must be >= 1")
35 if max_attempts < 1:
36 raise ValueError("max_attempts must be >= 1")
37 if retry_backoff_base_seconds < 0:
38 raise ValueError("retry_backoff_base_seconds must be >= 0")
40 self._queue = queue
41 self._handler = handler
42 self._worker_count = worker_count
43 self._lease_seconds = lease_seconds
44 self._max_attempts = max_attempts
45 self._retry_backoff_base_seconds = retry_backoff_base_seconds
46 self._job_types = job_types
47 self._source_locks: dict[str, asyncio.Lock] = {}
48 self._source_locks_guard = asyncio.Lock()
49 self._queue_io_guard = asyncio.Lock()
51 async def run_until_empty(self) -> int:
52 """Drain the queue once and return number of attempted jobs."""
53 processed_count = 0
54 processed_count_guard = asyncio.Lock()
55 active_workers: set[int] = set()
56 active_workers_guard = asyncio.Lock()
58 async def _set_worker_activity(worker_id: int, is_active: bool) -> int:
59 async with active_workers_guard:
60 if is_active:
61 active_workers.add(worker_id)
62 else:
63 active_workers.discard(worker_id)
64 return len(active_workers)
66 async def worker(worker_id: int) -> None:
67 nonlocal processed_count
69 while True:
70 async with self._queue_io_guard:
71 job = await self._queue.claim_next(
72 lease_seconds=self._lease_seconds,
73 job_types=self._job_types,
74 )
75 if job is None:
76 return
78 current_active_workers = await _set_worker_activity(worker_id, True)
80 payload, error_message = self._decode_payload(job.payload_json)
81 if error_message is not None:
82 async with self._queue_io_guard:
83 updated = await self._queue.mark_failed(
84 job.id, error_message, claim_attempt=job.attempts
85 )
86 if not updated:
87 logger.warning(
88 "job.claim_lost_on_terminal_transition",
89 job_id=job.id,
90 worker_id=worker_id,
91 active_workers=current_active_workers,
92 worker_count=self._worker_count,
93 )
94 async with processed_count_guard:
95 processed_count += 1
96 await _set_worker_activity(worker_id, False)
97 continue
99 try:
100 source_key = self._extract_source_key(payload)
101 except KeyError as exc:
102 async with self._queue_io_guard:
103 updated = await self._queue.mark_failed(
104 job.id, str(exc), claim_attempt=job.attempts
105 )
106 if not updated:
107 logger.warning(
108 "job.claim_lost_on_terminal_transition",
109 job_id=job.id,
110 worker_id=worker_id,
111 active_workers=current_active_workers,
112 worker_count=self._worker_count,
113 )
114 async with processed_count_guard:
115 processed_count += 1
116 await _set_worker_activity(worker_id, False)
117 continue
118 source_lock = await self._get_source_lock(source_key)
120 logger.info(
121 "job.claimed",
122 job_id=job.id,
123 job_type=job.type,
124 source_key=source_key,
125 attempt=job.attempts,
126 worker_id=worker_id,
127 active_workers=current_active_workers,
128 worker_count=self._worker_count,
129 )
131 async with source_lock:
132 logger.info(
133 "job.handler_started",
134 job_id=job.id,
135 job_type=job.type,
136 source_key=source_key,
137 attempt=job.attempts,
138 worker_id=worker_id,
139 active_workers=current_active_workers,
140 worker_count=self._worker_count,
141 )
142 t0 = time.monotonic()
143 claim_lost = False
144 renewal_interval = max(1, self._lease_seconds // 3)
145 # Mutable holder so _renew_lease closure can reference handler_task
146 # after it is created (asyncio is single-threaded; the first
147 # renewal sleep guarantees the holder is populated before use).
148 _handler_task_holder: list[asyncio.Task | None] = [None]
150 async def _renew_lease(
151 *,
152 current_job_id: int = job.id,
153 current_claim_attempt: int = job.attempts,
154 current_lease_seconds: int = self._lease_seconds,
155 current_renewal_interval: int = renewal_interval,
156 holder: list[asyncio.Task | None] = _handler_task_holder,
157 ) -> None:
158 nonlocal claim_lost
159 while True:
160 await asyncio.sleep(current_renewal_interval)
161 try:
162 async with self._queue_io_guard:
163 renewed = await self._queue.extend_visibility(
164 current_job_id,
165 current_lease_seconds,
166 claim_attempt=current_claim_attempt,
167 )
168 if not renewed:
169 # Claim was silently lost (job reclaimed or cancelled
170 # externally). Cancel the handler to abort side-effects.
171 logger.warning(
172 "job.claim_lost_on_renewal",
173 job_id=current_job_id,
174 )
175 claim_lost = True
176 if holder[0] is not None:
177 holder[0].cancel()
178 return
179 except Exception as exc:
180 logger.warning(
181 "job.lease_renew_failed",
182 job_id=current_job_id,
183 error=str(exc),
184 error_type=type(exc).__name__,
185 )
187 # Run handler as a task so the renewal loop can cancel it on claim loss.
188 handler_task = asyncio.create_task(self._handler(job.type, payload))
189 _handler_task_holder[0] = handler_task
190 renewal_task = asyncio.create_task(_renew_lease())
192 handler_exc: Exception | None = None
193 _external_cancelled = False
195 try:
196 await handler_task
197 except asyncio.CancelledError:
198 if not claim_lost:
199 # True external cancellation (not from claim-loss path).
200 # Ensure handler_task is stopped, then re-raise.
201 _external_cancelled = True
202 handler_task.cancel()
203 try:
204 await handler_task
205 except (asyncio.CancelledError, Exception):
206 pass
207 # If claim_lost: expected cancellation triggered by renewal;
208 # fall through to cleanup without re-raising.
209 except Exception as exc:
210 handler_exc = exc
211 finally:
212 renewal_task.cancel()
213 try:
214 await renewal_task
215 except asyncio.CancelledError:
216 pass
217 except Exception as exc:
218 logger.warning(
219 "job.lease_renew_teardown_failed",
220 job_id=job.id,
221 error=str(exc),
222 error_type=type(exc).__name__,
223 )
225 if _external_cancelled:
226 raise asyncio.CancelledError()
228 if claim_lost:
229 # Job was reclaimed by another worker; skip all status mutations
230 # to avoid overwriting the new owner's state.
231 logger.warning(
232 "job.claim_lost_skipping_update",
233 job_id=job.id,
234 job_type=job.type,
235 source_key=source_key,
236 attempt=job.attempts,
237 )
238 elif handler_exc is not None:
239 duration_ms = round((time.monotonic() - t0) * 1000)
240 # PermanentJobError always fails immediately, never retries
241 is_retry = (
242 not isinstance(handler_exc, PermanentJobError)
243 and job.attempts < self._max_attempts
244 )
245 retry_after_seconds = 0
246 if is_retry and self._retry_backoff_base_seconds > 0:
247 retry_after_seconds = self._retry_backoff_base_seconds * (
248 2 ** (job.attempts - 1)
249 )
250 async with self._queue_io_guard:
251 if is_retry:
252 updated = await self._queue.release_for_retry(
253 job.id,
254 str(handler_exc),
255 claim_attempt=job.attempts,
256 retry_after_seconds=retry_after_seconds,
257 )
258 else:
259 updated = await self._queue.mark_failed(
260 job.id,
261 str(handler_exc),
262 claim_attempt=job.attempts,
263 )
264 if updated:
265 if is_retry:
266 logger.info(
267 "job.retry_scheduled",
268 job_id=job.id,
269 job_type=job.type,
270 source_key=source_key,
271 attempt=job.attempts,
272 max_attempts=self._max_attempts,
273 retry_after_seconds=retry_after_seconds,
274 duration_ms=duration_ms,
275 error=str(handler_exc),
276 worker_id=worker_id,
277 active_workers=current_active_workers,
278 worker_count=self._worker_count,
279 )
280 else:
281 logger.info(
282 "job.failed",
283 job_id=job.id,
284 job_type=job.type,
285 source_key=source_key,
286 attempt=job.attempts,
287 max_attempts=self._max_attempts,
288 duration_ms=duration_ms,
289 error=str(handler_exc),
290 worker_id=worker_id,
291 active_workers=current_active_workers,
292 worker_count=self._worker_count,
293 )
294 else:
295 logger.warning(
296 "job.claim_lost_on_terminal_transition",
297 job_id=job.id,
298 job_type=job.type,
299 source_key=source_key,
300 attempt=job.attempts,
301 worker_id=worker_id,
302 active_workers=current_active_workers,
303 worker_count=self._worker_count,
304 )
305 else:
306 duration_ms = round((time.monotonic() - t0) * 1000)
307 async with self._queue_io_guard:
308 updated = await self._queue.mark_done(
309 job.id, claim_attempt=job.attempts
310 )
311 if updated:
312 logger.info(
313 "job.done",
314 job_id=job.id,
315 job_type=job.type,
316 source_key=source_key,
317 attempt=job.attempts,
318 duration_ms=duration_ms,
319 worker_id=worker_id,
320 active_workers=current_active_workers,
321 worker_count=self._worker_count,
322 )
323 else:
324 logger.warning(
325 "job.claim_lost_on_terminal_transition",
326 job_id=job.id,
327 job_type=job.type,
328 source_key=source_key,
329 attempt=job.attempts,
330 worker_id=worker_id,
331 active_workers=current_active_workers,
332 worker_count=self._worker_count,
333 )
335 async with processed_count_guard:
336 processed_count += 1
337 await _set_worker_activity(worker_id, False)
339 await asyncio.gather(
340 *(worker(worker_id) for worker_id in range(1, self._worker_count + 1))
341 )
342 return processed_count
344 async def _get_source_lock(self, source_key: str) -> asyncio.Lock:
345 async with self._source_locks_guard:
346 lock = self._source_locks.get(source_key)
347 if lock is None:
348 lock = asyncio.Lock()
349 self._source_locks[source_key] = lock
350 return lock
352 @staticmethod
353 def _decode_payload(payload_json: str) -> tuple[dict[str, Any], str | None]:
354 try:
355 payload = json.loads(payload_json)
356 except JSONDecodeError as exc:
357 return {}, f"Invalid payload_json: {exc.msg}"
359 if not isinstance(payload, dict):
360 return {}, "Invalid payload_json: expected JSON object"
361 return payload, None
363 @classmethod
364 def _extract_source_key(cls, payload: dict[str, Any]) -> str:
365 """Return the per-source concurrency key from the job payload.
367 Producers *must* set ``source_lock`` to a non-empty string. This is
368 the explicit contract: the pool will never silently fall back to a
369 global key, because that would serialize the entire pool and hide a
370 missing-field bug until production load.
372 Raises:
373 KeyError: if ``source_lock`` is absent or blank.
374 """
375 value = payload.get("source_lock")
376 if isinstance(value, str) and value.strip():
377 return value.strip()
378 raise KeyError(
379 "job payload is missing a non-empty 'source_lock' field — "
380 "all producers must set it explicitly"
381 )