diff --git a/CLAUDE.md b/CLAUDE.md index f54248011..2f726fe21 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -248,12 +248,36 @@ Per-user file quota enforcement tracked via the `file_count` and `file_quota` co **How it works:** - `files.created_by` records which user uploaded each file (nullable for pre-migration files) -- `users.file_count` is incremented/decremented in application code (in `PartitionFileManager`) — no SQL triggers -- Decrements use `func.greatest(file_count - N, 0)` to prevent negative values from race conditions +- `users.file_count` is incremented/decremented in application code — no SQL triggers +- Decrements use `GREATEST(file_count - N, 0)` to prevent negative values from race conditions - `delete_partition` queries per-uploader counts before cascade delete, then bulk decrements -- Quota check (`check_user_file_quota` in `openrag/api/dependencies/auth.py`) runs on upload, considering both indexed files and pending tasks -**Quota logic (`file_quota` column):** +**Atomic reserve/release (issue #664).** `file_count` is a **reserved + completed** +counter, not a "completed files" counter. Admission (`check_user_file_quota` in +`openrag/api/dependencies/auth.py`) charges a slot with one conditional UPDATE +(`UserRepository.try_reserve_file_slot`) so concurrent uploads cannot all read the +same pre-increment count and overshoot the quota. There is **no** completion-time +increment — `add_file_to_partition` only consumes the existing reservation. + +Consequences to respect when touching this code: +- The in-memory `TaskStateManager` pending count is **not** an admission input. It is + volatile (a restart zeroes it) and reserved uploads are already inside `file_count`. + Never add it back into a quota decision, and never add it to `file_count` when + reporting usage — that double-counts in-flight uploads. +- A reservation has an **owner**. Before dispatch it is the request's: the + `check_user_file_quota` yield-teardown releases it unless the router calls + `commit_quota_reservation(...)`. After dispatch it is the worker's: + `IndexerWorker.process_file` releases it in a `finally` unless the catalog write + reports a new row. Both quota-gated routes (`add_file`, `copy_file`) reserve, and + so do the MCP tools that create rows (`MCPService.index_url`, `MCPService.copy_file`), + which have no dependency chain and therefore reserve inline; + `put_file` (replace re-index) does not, since it reuses an existing row. +- Any new early return between admission and dispatch is automatically covered by the + teardown — but any new code path that *creates a file row without reserving*, or + *reserves without either committing or releasing*, leaks the counter. A leak is + silent and permanently narrows the user's quota. + +**Quota logic (`file_quota` column)** — the reserve SQL predicate reproduces exactly this: - `None` → use global default (`DEFAULT_FILE_QUOTA` env var, default `-1`) - `< 0` → unlimited - `>= 0` → specific limit diff --git a/openrag/api/dependencies/auth.py b/openrag/api/dependencies/auth.py index a0ab1ca2d..84d4b2324 100644 --- a/openrag/api/dependencies/auth.py +++ b/openrag/api/dependencies/auth.py @@ -254,40 +254,90 @@ def require_admin_or_self( ) +class QuotaReservation: + """One reserved file slot, owned by the request that admitted it. + + Issue #664: admission increments ``users.file_count`` *before* the + upload is dispatched, so the counter is now "reserved + completed" + rather than "completed". That makes the reservation a resource with an + owner — it must either be **consumed** (a ``files`` row is created for + it) or **released**. + + The handover point is dispatch: until the job is queued the request + owns the slot and :func:`check_user_file_quota`'s teardown releases it + on any error; once ``commit()`` is called the indexing worker owns it + and is responsible for releasing on failure/cancellation. So the + router must call ``commit()`` — via :func:`commit_quota_reservation` — + at exactly the moment responsibility transfers, and never earlier. + """ + + __slots__ = ("user_id", "committed") + + def __init__(self, user_id: int) -> None: + self.user_id = user_id + self.committed = False + + def commit(self) -> None: + """Hand the slot off to the worker; teardown will not release it.""" + self.committed = True + + +def commit_quota_reservation(reservation: object) -> None: + """Commit ``reservation`` when it is a real one; no-op otherwise. + + Routers receive whatever ``check_user_file_quota`` yields, and tests + routinely override that dependency with a plain stub. Guarding on the + type here keeps the routers free of ``if ... is not None`` noise and + keeps an overridden dependency from turning into an AttributeError. + """ + if isinstance(reservation, QuotaReservation): + reservation.commit() + + async def check_user_file_quota( user=Depends(current_user), auth_service=Depends(get_auth_service), - job_service=Depends(get_job_service), config=Depends(get_config), ): - """Check if user has reached their file quota.""" - default_file_quota = config.rdb.default_file_quota - if user.get("is_admin", False): - return user - user_quota = user.get("file_quota") - effective_quota = default_file_quota if user_quota is None else user_quota - if effective_quota < 0: - return user - + """Atomically reserve one file slot for this upload, or reject it. + + This is the admission gate for the two quota-bearing routes (``add_file`` + and ``copy_file``). It replaces the pre-#664 read-then-check, which + compared the request's stale ``file_count`` snapshot plus an in-memory + pending-task count against the quota and admitted every racer. + + The in-memory ``TaskStateManager`` count is deliberately **not** an + input any more: it is not durable (a restart zeroes it, reopening the + gate) and it is no longer needed — a reserved slot is already counted + in the durable ``file_count``. + + Yields a :class:`QuotaReservation`; on the way out, an uncommitted + reservation is released. That covers every path between admission and + dispatch — a 409 duplicate, a rejected/oversize upload, workspace + validation, a dispatch error, or a client disconnect — because FastAPI + propagates the endpoint's exception into this generator. + """ user_id = user.get("id") - pending_count = await job_service.get_user_pending_task_count(user_id) - - logger.debug( - "User file quota check", - user_id=user_id, - pending_count=pending_count, - ) + if user_id is None: + # No durable identity to charge (e.g. auth disabled). Nothing to + # reserve, so nothing can be enforced or leaked. + yield None + return try: - auth_service.validate_file_quota( - user, - pending_task_count=pending_count, - default_quota=default_file_quota, + new_count = await auth_service.reserve_file_slot( + user_id, + default_quota=config.rdb.default_file_quota, ) except OpenRAGError as exc: - raise HTTPException( - status_code=exc.status_code, - detail=exc.message, - ) from exc + logger.bind(user_id=user_id).info("Upload rejected: file quota exceeded.") + raise HTTPException(status_code=exc.status_code, detail=exc.message) from exc - return user + logger.bind(user_id=user_id, file_count=new_count).debug("Reserved a file slot.") + reservation = QuotaReservation(user_id) + try: + yield reservation + finally: + if not reservation.committed: + logger.bind(user_id=user_id).debug("Releasing an unconsumed file-slot reservation.") + await auth_service.release_file_slot(user_id) diff --git a/openrag/api/routers/admin/indexing.py b/openrag/api/routers/admin/indexing.py index 078075918..69d8fedf5 100644 --- a/openrag/api/routers/admin/indexing.py +++ b/openrag/api/routers/admin/indexing.py @@ -17,6 +17,7 @@ from api.dependencies.auth import ( check_user_file_quota, + commit_quota_reservation, current_user, current_user_partitions, ensure_partition_role, @@ -182,7 +183,13 @@ async def add_file( original_filename=original_filename, user=user, workspace_ids=parsed_workspace_ids, + quota_reserved=True, ) + # The job is queued: the worker now owns the reserved file slot and + # releases it if indexing fails or is cancelled. Committing before this + # point would leak the slot on a dispatch error; committing is skipped + # on every early return above, which releases it instead (#664). + commit_quota_reservation(_quota_check) return JSONResponse( status_code=status.HTTP_201_CREATED, @@ -400,7 +407,7 @@ async def copy_file_between_partitions( partition_service=partition_service, ) - await service.copy_file( + created = await service.copy_file( source_file_id=source_file_id, source_partition=source_partition, target_file_id=file_id, @@ -408,6 +415,12 @@ async def copy_file_between_partitions( metadata=metadata, user=user, ) + # Copy is synchronous, so the reservation is consumed here or not at all. + # ``created`` is False when no catalog row was written — the source had no + # chunks, or the target already existed (duplicate-at-catalog race) — in + # which case the slot goes back rather than leaking (#664). + if created: + commit_quota_reservation(_quota_check) return JSONResponse( status_code=status.HTTP_201_CREATED, content={"message": "File copied successfully."}, diff --git a/openrag/core/indexing/dispatcher.py b/openrag/core/indexing/dispatcher.py index 3250aaa61..f6aaeefe3 100644 --- a/openrag/core/indexing/dispatcher.py +++ b/openrag/core/indexing/dispatcher.py @@ -39,8 +39,14 @@ async def dispatch_indexing( embedder_name: str | None = None, require_existing_partition: bool = False, allow_legacy_require_existing_partition_retry: bool = False, + quota_reserved: bool = False, ) -> str: - """Queue an (re)indexing job, register its task state, return its id.""" + """Queue an (re)indexing job, register its task state, return its id. + + ``quota_reserved`` marks that one ``users.file_count`` slot was + charged at admission (#664) and now belongs to the job: the worker + releases it if the file never reaches the catalog. + """ ... @abstractmethod @@ -66,8 +72,13 @@ async def copy_file( metadata: dict, partition: str, user: dict | None, - ) -> None: - """Copy a file's chunks into another partition / file id.""" + ) -> bool: + """Copy a file's chunks into another partition / file id. + + Returns whether a catalog row was actually created (False for an + empty source or an already-existing target) — the caller needs it to + settle the reserved quota slot (#664). + """ ... @abstractmethod diff --git a/openrag/core/models/catalog.py b/openrag/core/models/catalog.py index c650c5049..9c225d113 100644 --- a/openrag/core/models/catalog.py +++ b/openrag/core/models/catalog.py @@ -53,12 +53,32 @@ class DocumentRecord(BaseModel): class IndexationJob(BaseModel): - """An indexation job tracking batch document processing.""" + """A durable indexation job — one row per dispatched indexing task. + + The record mirrors what the in-memory ``TaskStateManager`` Ray actor holds, + but survives restarts and is operator-visible (issue #660). ``id`` is the + dispatcher's ``task_id``, so the durable row and the hot-cache entry share + one identity. + + ``status`` reuses :class:`DocumentStatus` rather than :class:`JobStatus` + because a job here tracks exactly one file and must reproduce the actor's + state taxonomy verbatim (``QUEUED`` -> ``SERIALIZING`` -> ... -> + ``COMPLETED`` / ``FAILED`` / ``CANCELLED``); :class:`JobStatus` is the + coarser roll-up reserved for a future batch-level job entity. + + ``error`` is bounded at write time (see + :func:`core.utils.text.truncate_error_text`) — a full traceback is + unbounded input and this row is retained. + """ id: str = Field(default_factory=lambda: str(uuid.uuid4())) - status: JobStatus = JobStatus.QUEUED - total_documents: int = 0 + status: DocumentStatus = DocumentStatus.QUEUED partition: str = "default" + file_id: str | None = None + user_id: int | None = None + job_metadata: dict[str, Any] = Field(default_factory=dict) + error: str | None = None created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) started_at: datetime | None = None completed_at: datetime | None = None + updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/openrag/core/ports/job_repo.py b/openrag/core/ports/job_repo.py index 5e20b6830..1dc42d802 100644 --- a/openrag/core/ports/job_repo.py +++ b/openrag/core/ports/job_repo.py @@ -3,10 +3,17 @@ from __future__ import annotations from abc import ABC, abstractmethod +from datetime import datetime from typing import Any from openrag.core.models.catalog import IndexationJob +# The non-terminal states an indexation job passes through. Shared by the +# repository implementations and the read models so "active"/"pending" means +# exactly one thing across the stack. +ACTIVE_JOB_STATES: tuple[str, ...] = ("QUEUED", "SERIALIZING", "CHUNKING", "INSERTING") +TERMINAL_JOB_STATES: tuple[str, ...] = ("COMPLETED", "FAILED", "CANCELLED") + class JobRepository(ABC): """CRUD operations for indexation jobs.""" @@ -18,7 +25,51 @@ async def create_job(self, job: IndexationJob) -> IndexationJob: ... async def get_job(self, job_id: str) -> IndexationJob | None: ... @abstractmethod - async def list_jobs(self, status: str | None = None, offset: int = 0, limit: int = 50) -> list[IndexationJob]: ... + async def list_jobs( + self, + status: str | None = None, + offset: int = 0, + limit: int = 50, + user_id: int | None = None, + ) -> list[IndexationJob]: + """Return jobs newest-first. + + ``status`` accepts an exact state (case-insensitive) or the pseudo-status + ``"active"``, which expands to :data:`ACTIVE_JOB_STATES`. ``user_id`` + scopes the result to one uploader; ``None`` means every job. + """ @abstractmethod async def update_job(self, job_id: str, **fields: Any) -> IndexationJob | None: ... + + @abstractmethod + async def mark_failed_if_not_cancelled(self, job_id: str, *, error: str, completed_at: datetime) -> bool: + """Record a FAILED outcome unless the row is already CANCELLED. + + Arbitration lives here rather than in the in-memory task actor because + the durable row is the only participant guaranteed to still exist. An + actor that restarted or evicted the entry cannot say whether the user + cancelled, and treating "I don't know" as "cancelled" leaves the job in a + non-terminal state that retention never sweeps — a permanent phantom in + the queue views, which is the failure #660 exists to remove. + + Returns ``True`` if the row moved to FAILED. + """ + + @abstractmethod + async def count_by_status(self) -> dict[str, int]: + """Return ``{status: count}`` over every retained job. + + Cheaper than paging the whole table for the queue-info roll-up. + """ + + @abstractmethod + async def purge_terminal_jobs(self, *, older_than_seconds: int, keep_last: int) -> int: + """Evict terminal jobs, returning how many rows were removed. + + Retention is what keeps the durable store from repeating the unbounded + growth of the in-memory actor it replaced. Both bounds apply: a terminal + job is removed once it is older than ``older_than_seconds`` **or** once + it falls outside the ``keep_last`` most recently completed rows. + In-flight jobs are never touched. + """ diff --git a/openrag/core/ports/user_repo.py b/openrag/core/ports/user_repo.py index 747eadf0c..ff3d68f87 100644 --- a/openrag/core/ports/user_repo.py +++ b/openrag/core/ports/user_repo.py @@ -51,6 +51,30 @@ async def delete_user(self, user_id: int) -> bool: ... @abstractmethod async def count_users(self) -> int: ... + # ── File-quota reserve / release ────────────────────────────────── + # + # ``users.file_count`` is a *reserved + completed* counter (issue #664). + # Admission reserves a slot with a single conditional UPDATE so N + # concurrent uploads can never overshoot the quota; whoever holds an + # unconsumed reservation must release it. + + @abstractmethod + async def try_reserve_file_slot(self, user_id: int, *, default_quota: int) -> int | None: + """Atomically claim one file slot against the user's quota. + + Returns the post-increment ``file_count`` when the slot was + granted, or ``None`` when the user is at (or over) quota — or does + not exist. Quota semantics: admins bypass; a ``NULL`` per-user + ``file_quota`` falls back to ``default_quota``; a *resolved* quota + ``< 0`` means unlimited. + """ + ... + + @abstractmethod + async def release_file_slot(self, user_id: int) -> None: + """Give back one reserved slot (clamped at zero). Idempotent-ish.""" + ... + # ── API keys ────────────────────────────────────────────────────── @abstractmethod diff --git a/openrag/core/utils/text.py b/openrag/core/utils/text.py index 17aa0d435..85e64066c 100644 --- a/openrag/core/utils/text.py +++ b/openrag/core/utils/text.py @@ -14,6 +14,15 @@ DEFAULT_FALLBACK_ENCODING = "utf-8" +# Upper bound for a stored error/traceback string (issue #660). Tracebacks are +# unbounded input: they were kept in full both in the detached TaskStateManager +# actor (which never evicted) and, now, in a retained ``jobs`` row. The tail is +# what is kept — the exception type/message and the innermost frames are the +# diagnostic payload; the outermost frames are the same dispatcher stack on +# every task. +MAX_ERROR_TEXT_CHARS = 4000 + + logger = get_logger() @@ -44,6 +53,33 @@ def get_num_tokens(): return _cached_length_function +def truncate_error_text(text: str | None, max_chars: int = MAX_ERROR_TEXT_CHARS) -> str | None: + """Bound an error/traceback string, keeping its most diagnostic tail. + + Returns ``text`` unchanged when it already fits (the overwhelmingly common + case), so this is free on the happy path. Oversized text is replaced by a + marker line naming the number of dropped characters followed by the last + ``max_chars`` characters, which makes the truncation auditable rather than + silent. + + ``max_chars`` bounds the retained *original* text, not the returned string: + the audit marker is overhead on top, so the result runs ~35 characters + longer. The point is to stop unbounded growth, and the caller's column is + unbounded ``VARCHAR`` — a hard ceiling on the return value would have to + either eat into the tail (the exception message lives at the very end, and + is the whole reason the tail is what we keep) or iterate to a fixed point + against the marker's own digit count. Neither is worth it here. + """ + if text is None: + return None + if max_chars <= 0: + raise ValueError("max_chars must be positive") + if len(text) <= max_chars: + return text + dropped = len(text) - max_chars + return f"[truncated {dropped} of {len(text)} chars]\n...{text[-max_chars:]}" + + def decode_bytes(raw: bytes, encoding: str | None = None) -> str: """Decode ``raw`` to ``str`` with a UTF-8-first detection strategy. diff --git a/openrag/di/container.py b/openrag/di/container.py index c2d8e11fa..1b0a3cf00 100644 --- a/openrag/di/container.py +++ b/openrag/di/container.py @@ -583,6 +583,7 @@ def indexing_service(self) -> IndexingService: document_repo=self.document_repo, workspace_repo=self.workspace_repo, collection=settings.vectordb.collection_name, + job_repo=self.job_repo, ), config=settings, partition_service=self.partition_service, @@ -593,15 +594,18 @@ def indexing_service(self) -> IndexingService: def job_service(self) -> JobService: """JobService — lazily built, cached for the container's lifetime. - Wraps the ``TaskStateManager`` Ray actor directly (8H excepts - JobService); resolved lazily so the actor only needs to exist at - first request. + Reads the durable ``jobs`` table, with the ``TaskStateManager`` Ray + actor as the fallback (8H excepts JobService for wrapping the actor); + resolved lazily so the actor only needs to exist at first request. """ if self._job_service is None: from services.orchestrators.job_service import JobService from services.workers.bootstrap import get_task_state_manager - self._job_service = JobService(task_state_manager=get_task_state_manager()) + self._job_service = JobService( + task_state_manager=get_task_state_manager(), + job_repo=self.job_repo, + ) return self._job_service @property @@ -645,8 +649,10 @@ def mcp_service(self) -> MCPService: indexing_service=self.indexing_service, job_service=self.job_service, conversion_service=self.conversion_service, + auth_service=self.auth_service, vector_store=self.vector_store, collection=settings.vectordb.collection_name, + default_file_quota=settings.rdb.default_file_quota, default_top_k=mcp_cfg.default_top_k, max_top_k=mcp_cfg.max_top_k, similarity_threshold=mcp_cfg.similarity_threshold, diff --git a/openrag/services/orchestrators/auth_service.py b/openrag/services/orchestrators/auth_service.py index be240a047..95740aca9 100644 --- a/openrag/services/orchestrators/auth_service.py +++ b/openrag/services/orchestrators/auth_service.py @@ -559,6 +559,53 @@ def authorize(cls, *, user: Any, action: str, resource: dict[str, Any] | None = case _: raise ValueError(f"Unknown authorization action: {action!r}") + # ── File quota ──────────────────────────────────────────────────── + + async def reserve_file_slot(self, user_id: int, *, default_quota: int) -> int: + """Atomically admit one upload against the user's quota (issue #664). + + Admission used to be read-then-check: every concurrent request read + the same pre-increment ``file_count`` (plus a volatile in-memory + pending count) and all of them passed. Reserving DB-side collapses + the check and the admit into one statement, so a burst of ``K`` + uploads at quota ``N`` admits exactly the slots that are free. + + Returns the post-reserve ``file_count``. Raises + :class:`~core.utils.exceptions.OpenRAGError` (``FILE_QUOTA_EXCEEDED``, + 403) when no slot is available — the same code/status the previous + :meth:`validate_file_quota` raised, so the API contract is unchanged. + + The caller **owns** the returned reservation: it must either be + consumed (a ``files`` row gets created for it) or handed back via + :meth:`release_file_slot`. + """ + new_count = await self._user_repo.try_reserve_file_slot(user_id, default_quota=default_quota) + if new_count is None: + raise OpenRAGError( + "File quota exceeded. Delete a file or ask an administrator to raise your quota.", + code="FILE_QUOTA_EXCEEDED", + status_code=403, + ) + return new_count + + async def release_file_slot(self, user_id: int) -> None: + """Hand a reserved-but-unused slot back (never raises). + + Release runs on cleanup paths (upload rejected after admission, + indexing failed, task cancelled), where a raising cleanup would mask + the original error. A *failed* release only costs the user one slot + until an admin fixes the count; a raised exception here would lose + the real failure, so we log loudly and swallow. + """ + try: + await self._user_repo.release_file_slot(user_id) + except Exception as exc: # noqa: BLE001 — cleanup must never mask the real error + logger.bind(user_id=user_id).error( + "Failed to release a reserved file slot; the user's file_count is now " + "one too high and must be reconciled manually.", + error=str(exc), + ) + @classmethod def validate_file_quota( cls, @@ -569,6 +616,13 @@ def validate_file_quota( ) -> None: """Pure quota check (the pending-task count is supplied by the caller). + .. deprecated:: issue #664 + Superseded by :meth:`reserve_file_slot` for *admission* — this + read-then-check shape is exactly the TOCTOU race #664 fixes and + must not gate uploads again. Retained only as the pure-function + expression of the quota rules (and for read-only callers that + want to report quota state without reserving). + Quota semantics (shared with ``check_user_file_quota``): admins bypass; a per-user ``file_quota`` of ``None`` falls back to the global default; a *resolved* quota ``< 0`` means unlimited. So a diff --git a/openrag/services/orchestrators/indexing_service.py b/openrag/services/orchestrators/indexing_service.py index cf8f9f6c8..d1fe6f080 100644 --- a/openrag/services/orchestrators/indexing_service.py +++ b/openrag/services/orchestrators/indexing_service.py @@ -199,11 +199,18 @@ async def add_file( user: dict | None, workspace_ids: list[str] | None = None, replace: bool = False, + quota_reserved: bool = False, ) -> str: """Assemble metadata and queue an (re)indexing job; return its task id. Workspace association happens inside the worker's ``add_file`` after a successful index — the router only pre-validates the ids. + + ``quota_reserved`` tells the worker that the caller already charged + one slot to ``users.file_count`` at admission (#664), so the worker + must give it back if the file never lands in the catalog. It is + False for ``replace`` re-indexing, which reuses an existing row and + therefore never reserves. """ full_metadata = self._build_metadata( metadata=metadata, @@ -228,6 +235,7 @@ async def add_file( embedder_name=embedder_name, require_existing_partition=require_existing_partition, allow_legacy_require_existing_partition_retry=legacy_actor_preserves_partition_guard, + quota_reserved=quota_reserved, ) async def delete_file(self, file_id: str, partition: str) -> None: @@ -253,11 +261,18 @@ async def copy_file( target_partition: str, metadata: dict, user: dict | None, - ) -> None: + ) -> bool: + """Copy a file into another partition; return whether a row was created. + + The boolean is what the router needs to decide the fate of the quota + slot it reserved at admission (#664): a copy that writes no catalog + row (empty source, or the target file already exists) consumed no + slot and must not keep one. + """ metadata = dict(metadata or {}) metadata["file_id"] = target_file_id metadata["partition"] = target_partition - await self._dispatcher.copy_file(source_file_id, metadata, source_partition, user) + return bool(await self._dispatcher.copy_file(source_file_id, metadata, source_partition, user)) # ------------------------------------------------------------------ # Task state diff --git a/openrag/services/orchestrators/job_service.py b/openrag/services/orchestrators/job_service.py index 4ca4978ef..c4a7b4678 100644 --- a/openrag/services/orchestrators/job_service.py +++ b/openrag/services/orchestrators/job_service.py @@ -1,15 +1,21 @@ """JobService — task-queue queries (Phase 8D.2). -Thin wrapper around the ``TaskStateManager`` Ray actor, extracted from -``routers/queue.py``. Aggregation/filtering (the active-status rollup, -the per-status counts, the ``?task_status=`` filter) is business logic -and lives here; ``request.url_for`` link building stays in the thin -router (HTTP transport). +Reads job state from the durable ``jobs`` table (``JobRepository``) and falls +back to the ``TaskStateManager`` Ray actor. Aggregation/filtering (the +active-status rollup, the per-status counts, the ``?task_status=`` filter) is +business logic and lives here; ``request.url_for`` link building stays in the +thin router (HTTP transport). + +Issue #660 inverted the roles: Postgres is the source of truth and the actor is +a hot cache that evicts settled tasks and is wiped by a restart, so reading the +actor first would make completed work vanish from the queue views. The actor is +kept as the fallback for the two cases where the durable store cannot answer — +it is not wired (no job repository) or it is unreachable — because a degraded, +restart-local view of the queue beats a 500. This is the one orchestrator that legitimately keeps Ray remote calls during the shim — 8H verification explicitly excepts JobService -wrapping ``TaskStateManager``. Phase 9 swaps the actor for a DB-backed -job repository (this service is the hook point for that P0 feature). +wrapping ``TaskStateManager``. """ from __future__ import annotations @@ -17,15 +23,30 @@ from collections import Counter from typing import Any -_ACTIVE_STATES = ("QUEUED", "SERIALIZING", "CHUNKING", "INSERTING") +from core.ports.job_repo import ACTIVE_JOB_STATES as _ACTIVE_STATES +from core.utils.logging import get_logger + +logger = get_logger() + +# Upper bound on a single ``list_tasks`` page against the durable store. The +# route is unpaginated (it used to read an in-memory dict), so without this a +# deployment with a full retention window would try to serialize every retained +# job into one response. +# +# Retention keeps up to ``JOB_RETENTION_MAX_ROWS`` (10k) rows, so this cap is +# reachable. ``list_tasks`` asks for one row more than it will return, which is +# what lets it tell "exactly _LIST_LIMIT jobs" apart from "more than we will +# show" and say so, rather than handing back a short answer that looks complete. +_LIST_LIMIT = 1000 class JobService: - """Queue/worker introspection over the TaskStateManager actor.""" + """Queue/worker introspection over the durable job store.""" - def __init__(self, task_state_manager: Any, timeout: float = 60.0) -> None: + def __init__(self, task_state_manager: Any, timeout: float = 60.0, job_repo: Any = None) -> None: self._tsm = task_state_manager self._timeout = timeout + self._job_repo = job_repo async def _call(self, future: Any, task_description: str) -> Any: """Route TaskStateManager calls through the centralized Ray helper. @@ -43,6 +64,46 @@ async def _call(self, future: Any, task_description: str) -> Any: task_description=task_description, ) + async def _from_jobs(self, operation: str, call: Any) -> Any: + """Run a durable read, or return ``None`` to signal "use the cache". + + Any repository failure is a fallback trigger, not a request failure: the + actor still holds recent state, and a queue view is a diagnostic surface + — the moment it 500s is exactly the moment it is being looked at. + + ``None`` means "the durable read did not happen". It does **not** mean + "no rows": the aggregate readers treat an *empty* result as a miss too + and fall through to the cache on their own, because the table starts + empty while the detached actor is already holding live tasks. The + single-row reader (:meth:`get_task_details`) keeps the ``is not None`` + test — there, ``None`` already is the miss. + """ + if self._job_repo is None: + return None + try: + return await call() + except Exception as exc: # noqa: BLE001 - fall back to the in-memory cache + logger.warning( + "Durable job read failed; falling back to the in-memory task cache", + operation=operation, + error=str(exc), + ) + return None + + @staticmethod + def _job_details(job: Any) -> dict: + """Project a durable job onto the actor's ``details`` shape. + + The two read paths must be indistinguishable to callers — the API + contract (and ``require_task_owner``) predates the durable store. + """ + return { + "file_id": job.file_id, + "partition": job.partition, + "metadata": job.job_metadata, + "user_id": job.user_id, + } + @staticmethod def _format_pool_info(worker_info: dict[str, int]) -> dict[str, int]: """Condense ``SerializerQueue.pool_info()`` into the API shape.""" @@ -53,8 +114,15 @@ def _format_pool_info(worker_info: dict[str, int]) -> dict[str, int]: } async def get_queue_info(self) -> dict: - all_states: dict = await self._call(self._tsm.get_all_states.remote(), "get_all_states") - status_counts = Counter(all_states.values()) + status_counts = await self._from_jobs("count_by_status", lambda: self._job_repo.count_by_status()) + if not status_counts: + # Empty, not just failed: an unpopulated ``jobs`` table is a durable + # *miss*, not an authoritative "nothing is running". The actor is + # detached, so it outlives the API restart that first deploys this — + # every task dispatched before the cutover has no row, and reporting + # zero active while workers are indexing is worse than a stale count. + all_states: dict = await self._call(self._tsm.get_all_states.remote(), "get_all_states") + status_counts = Counter(all_states.values()) active = {s: status_counts.get(s, 0) for s in _ACTIVE_STATES} task_summary = { @@ -82,8 +150,44 @@ async def list_tasks( - any other value → exact match (case-insensitive) - ``None`` → all tasks + Capped at ``_LIST_LIMIT`` rows; hitting the cap logs a warning, since the + route has no way to signal a partial answer in its response body. + The router decorates each row with the status / error URLs. """ + if not is_admin and user_id is None: + # Fail closed. list_jobs(user_id=None) means "every job", so an + # anonymous non-admin would otherwise be handed the whole table. + # Unreachable today (the HTTP path always resolves an id), but the + # unset default of the MCP _USER_ID ContextVar is exactly None, + # so the escalating value is one wiring change away. + return [] + jobs = await self._from_jobs( + "list_jobs", + lambda: self._job_repo.list_jobs( + status=task_status, + limit=_LIST_LIMIT + 1, # the extra row is the truncation probe + user_id=None if is_admin else user_id, + ), + ) + if jobs: + if len(jobs) > _LIST_LIMIT: + jobs = jobs[:_LIST_LIMIT] + logger.warning( + "Task list truncated at the page cap; the response is not the whole queue", + limit=_LIST_LIMIT, + task_status=task_status, + user_id=None if is_admin else user_id, + ) + # Filtering happened in SQL; the fallback below has to do it itself. + return [{"task_id": j.id, "state": j.status.value, "details": self._job_details(j)} for j in jobs] + + # No rows is a durable *miss*, not proof of an empty queue — fall through + # to the cache rather than answer ``[]`` authoritatively. See the note in + # ``get_queue_info``: the detached actor holds tasks the table never got. + # The cost is one actor call on a genuinely-empty query, which is exactly + # what this route did before the durable store existed. + if is_admin: all_info: dict[str, dict] = await self._call(self._tsm.get_all_info.remote(), "get_all_info") else: @@ -102,8 +206,19 @@ async def list_tasks( async def get_user_pending_task_count(self, user_id: int | None) -> int: """Pending (not-yet-completed) indexing tasks for one user. - Used by UserService for the quota-usage block of ``/users/info`` - (the legacy router called the actor directly from the handler). + Purely informational: UserService reports it as ``pending_files`` in the + quota-usage block of ``/users/info``. It is **not** a correctness input + anywhere — since #664 admission reserves a slot in ``users.file_count`` + directly, so an in-flight upload is already charged and adding this on + top would double-count it. Do not reintroduce it into a quota decision. + + Deliberately *not* served from the durable store, unlike every other read + here. A job row only leaves the active states when a worker writes a + terminal transition, so a job orphaned by a crash would be reported as + pending forever (retention only sweeps terminal rows). The in-memory + count is wrong in the opposite, self-healing direction: a restart clears + it. Serving this from Postgres is safe once orphaned in-flight jobs are + reconciled at startup — tracked in #676. """ return await self._call( self._tsm.get_user_pending_task_count.remote(user_id), @@ -112,6 +227,9 @@ async def get_user_pending_task_count(self, user_id: int | None) -> int: async def get_task_details(self, task_id: str) -> dict | None: """Return task details for ownership checks and status routes.""" + job = await self._from_jobs("get_job", lambda: self._job_repo.get_job(task_id)) + if job is not None: + return self._job_details(job) return await self._call( self._tsm.get_details.remote(task_id), f"get_details({task_id})", diff --git a/openrag/services/orchestrators/mcp_service.py b/openrag/services/orchestrators/mcp_service.py index 326558348..4e598634e 100644 --- a/openrag/services/orchestrators/mcp_service.py +++ b/openrag/services/orchestrators/mcp_service.py @@ -42,6 +42,7 @@ if TYPE_CHECKING: from core.vector_stores import VectorStore + from services.orchestrators.auth_service import AuthService from services.orchestrators.conversion_service import ConversionService from services.orchestrators.indexing_service import IndexingService from services.orchestrators.job_service import JobService @@ -88,8 +89,10 @@ def __init__( indexing_service: IndexingService, job_service: JobService, conversion_service: ConversionService, + auth_service: AuthService, vector_store: VectorStore, collection: str, + default_file_quota: int, default_top_k: int = 5, max_top_k: int = 50, similarity_threshold: float = 0.8, @@ -101,8 +104,12 @@ def __init__( self._indexing = indexing_service self._jobs = job_service self._conversion = conversion_service + # Required, not optional: these tools create ``files`` rows, and a row + # created without a reservation silently bypasses the file quota (#664). + self._auth = auth_service self._vector_store = vector_store self._collection = collection + self._default_file_quota = default_file_quota self.default_top_k = default_top_k self.max_top_k = max_top_k self.similarity_threshold = similarity_threshold @@ -557,14 +564,25 @@ async def copy_file( if await self._partitions.file_exists(dest_file_id, dest_partition): raise FileExistsError(f"File '{dest_file_id}' already exists in partition '{dest_partition}'") - await self._indexing.copy_file( - source_file_id=source_file_id, - source_partition=source_partition, - target_file_id=dest_file_id, - target_partition=dest_partition, - metadata=_strip_protected_metadata(extra_metadata), - user={"id": user_id} if user_id is not None else None, - ) + # Admission gate (#664) — same reasoning as ``index_url``. The copy is + # synchronous, so the slot is consumed here or handed straight back: + # ``created`` is False for an empty source or an already-existing target. + reserved_user_id = user_id + if reserved_user_id is not None: + await self._auth.reserve_file_slot(reserved_user_id, default_quota=self._default_file_quota) + created = False + try: + created = await self._indexing.copy_file( + source_file_id=source_file_id, + source_partition=source_partition, + target_file_id=dest_file_id, + target_partition=dest_partition, + metadata=_strip_protected_metadata(extra_metadata), + user={"id": user_id} if user_id is not None else None, + ) + finally: + if reserved_user_id is not None and not created: + await self._auth.release_file_slot(reserved_user_id) return { "source_partition": source_partition, "source_file_id": source_file_id, @@ -605,31 +623,51 @@ async def index_url( if await self._partitions.file_exists(file_id, partition): raise FileExistsError(f"File '{file_id}' already exists in partition '{partition}'") - filename = Path(urlparse(url).path.rstrip("/")).name or file_id - suffix = Path(filename).suffix or "" - with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: - tmp_path = Path(tmp.name) + # Admission gate (#664). ``users.file_count`` is a reserved+completed + # counter and there is no completion-time increment any more, so a path + # that creates a ``files`` row without reserving is invisible to the + # quota — while the delete path still decrements it, driving the count + # *below* reality and handing out free slots. The HTTP ``add_file`` route + # reserves in a FastAPI dependency; this tool has no dependency chain and + # must reserve itself. Done before the download so an over-quota call + # costs no bandwidth. + reserved_user_id = user_id + if reserved_user_id is not None: + await self._auth.reserve_file_slot(reserved_user_id, default_quota=self._default_file_quota) + dispatched = False try: - await self._safe_download(url, tmp_path) - except Exception as exc: - tmp_path.unlink(missing_ok=True) - raise RuntimeError(f"Failed to download '{url}': {exc}") from exc - - metadata = _strip_protected_metadata(extra_metadata) - metadata["source_url"] = url - guessed_mime, _ = mimetypes.guess_type(filename) - if guessed_mime and "mimetype" not in metadata: - metadata["mimetype"] = guessed_mime - - task_id = await self._indexing.add_file( - file_path=str(tmp_path), - file_id=file_id, - partition=partition, - metadata=metadata, - sanitized_filename=filename, - original_filename=filename, - user={"id": user_id, "is_admin": is_admin} if user_id is not None else None, - ) + filename = Path(urlparse(url).path.rstrip("/")).name or file_id + suffix = Path(filename).suffix or "" + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: + tmp_path = Path(tmp.name) + try: + await self._safe_download(url, tmp_path) + except Exception as exc: + tmp_path.unlink(missing_ok=True) + raise RuntimeError(f"Failed to download '{url}': {exc}") from exc + + metadata = _strip_protected_metadata(extra_metadata) + metadata["source_url"] = url + guessed_mime, _ = mimetypes.guess_type(filename) + if guessed_mime and "mimetype" not in metadata: + metadata["mimetype"] = guessed_mime + + task_id = await self._indexing.add_file( + file_path=str(tmp_path), + file_id=file_id, + partition=partition, + metadata=metadata, + sanitized_filename=filename, + original_filename=filename, + user={"id": user_id, "is_admin": is_admin} if user_id is not None else None, + quota_reserved=reserved_user_id is not None, + ) + # Queued: the worker owns the slot from here and releases it if the + # file never reaches the catalog. + dispatched = True + finally: + if reserved_user_id is not None and not dispatched: + await self._auth.release_file_slot(reserved_user_id) return { "partition": partition, "file_id": file_id, diff --git a/openrag/services/orchestrators/user_service.py b/openrag/services/orchestrators/user_service.py index 7b289ee74..3a2d7bb4b 100644 --- a/openrag/services/orchestrators/user_service.py +++ b/openrag/services/orchestrators/user_service.py @@ -134,6 +134,14 @@ async def get_current_user_info(self, user: dict) -> dict: a negative *global default* only makes users unlimited when they have no per-user override). ``file_quota`` is surfaced as ``-1`` when unlimited. + + Since #664 ``file_count`` is a *reserved + completed* counter: a + slot is charged at admission, so an upload that is still indexing is + already inside ``file_count``. ``total_files`` is therefore just + ``file_count`` — adding the in-memory pending count on top, as this + did before, would double-count every in-flight upload and report a + usage the quota gate does not actually enforce. ``pending_files`` + stays as informational "how many of those are still indexing". """ is_admin = user.get("is_admin", False) if is_admin: @@ -146,14 +154,14 @@ async def get_current_user_info(self, user: dict) -> dict: user_quota = float("inf") file_count = user.get("file_count", 0) + # Informational only — never a correctness input (see #664). pending_count = await self._job_service.get_user_pending_task_count(user.get("id")) - total = file_count + pending_count return { **user, "file_count": file_count, "pending_files": pending_count, - "total_files": total, + "total_files": file_count, "file_quota": -1 if user_quota == float("inf") else user_quota, } diff --git a/openrag/services/persistence/__init__.py b/openrag/services/persistence/__init__.py index 30477f4b6..db07f80a3 100644 --- a/openrag/services/persistence/__init__.py +++ b/openrag/services/persistence/__init__.py @@ -11,9 +11,11 @@ ``PgDocumentRepository``, ``PgUserRepository``, ``PgPartitionRepository``, ``PgPartitionMembershipRepository``, ``PgOIDCSessionRepository``, ``PgWorkspaceRepository``. + - Real (durable indexation job state, issue #660): + ``PgJobRepository``. - Stubs (post-refactoring features — raise :class:`StubRepositoryError` on every call): - ``PgJobRepository``, ``PgChunkRepository``, + ``PgChunkRepository``, ``PgPromptRepository``, ``PgConversationRepository``, ``PgAuditLogRepository``, ``PgIdempotencyRepository``, ``PgEntityRepository``, ``PgTopicTagRepository``, diff --git a/openrag/services/persistence/document_repo.py b/openrag/services/persistence/document_repo.py index f8fcf35d4..309c4afba 100644 --- a/openrag/services/persistence/document_repo.py +++ b/openrag/services/persistence/document_repo.py @@ -279,7 +279,9 @@ async def add_file_to_partition( # noqa: PLR0913 — legacy signature pinned """TODO(phase-9): remove. Mirror of legacy ``add_file_to_partition``. Creates the partition row on first use (legacy behaviour). Returns - ``False`` if a row with the same (file_id, partition) already exists. + ``False`` if a row with the same (file_id, partition) already exists — + callers holding a quota reservation must release it on ``False``, + since no file was created for it (#664). ``indexed_at`` pins the indexation timestamp so it matches the Milvus chunks; when ``None`` the ``files.indexed_at`` server default applies. @@ -356,11 +358,13 @@ async def add_file_to_partition( # noqa: PLR0913 — legacy signature pinned f"INSERT INTO files ({', '.join(columns)}) VALUES ({placeholders})", *values, ) - if user_id is not None: - await conn.execute( - "UPDATE users SET file_count = file_count + 1 WHERE id = $1", - user_id, - ) + # No ``file_count`` increment here. Since #664 the uploader's + # slot is charged atomically at *admission* + # (``UserRepository.try_reserve_file_slot``) and this insert + # merely consumes that reservation — incrementing again would + # double-count. The paths that never reach this insert release + # the reservation instead; see ``check_user_file_quota`` and + # ``IndexerWorker.process_file``. return True async def remove_file_from_partition(self, file_id: str, partition: str) -> bool: diff --git a/openrag/services/persistence/job_repo.py b/openrag/services/persistence/job_repo.py index 55cbbf3ce..ab0ccb711 100644 --- a/openrag/services/persistence/job_repo.py +++ b/openrag/services/persistence/job_repo.py @@ -1,41 +1,319 @@ -"""Stub :class:`JobRepository` — see ``_stubs.py`` for the rationale. - -Job state is currently tracked in-memory by the -:class:`components.indexer.indexer.TaskStateManager` Ray actor. The -post-refactoring P0 feature is to persist jobs to Postgres so they -survive restarts and become visible to operators. When that lands, -swap the body of each method for an asyncpg implementation against a -new ``jobs`` table — the port shape is already pinned by Phase 4. +"""asyncpg-backed :class:`JobRepository` — durable indexation job state. + +Replaces the Phase 7A.2 stub (issue #660). Job state used to live *only* in the +detached ``TaskStateManager`` Ray actor: unbounded, volatile and invisible to +operators, so any restart mid-batch made the in-flight work unobservable and +un-cancellable. The ``jobs`` table (one row per dispatched task, see +``schema.py``) is now the source of truth; the actor is a hot cache in front of +it. + +Two invariants shape this module: + +* **Bounded rows.** Terminal jobs are evicted by :meth:`purge_terminal_jobs` + and stored errors are truncated at write time — otherwise the durable store + would just reproduce the in-memory leak on disk. +* **Column allowlisting.** :meth:`update_job` takes ``**fields`` from callers + in the worker path; only known columns reach the SQL string. """ from __future__ import annotations -from typing import Any +from collections.abc import Callable +from typing import TYPE_CHECKING, Any from core.models.catalog import IndexationJob -from core.ports.job_repo import JobRepository -from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented +from core.ports.job_repo import ACTIVE_JOB_STATES, TERMINAL_JOB_STATES, JobRepository +from core.utils.text import truncate_error_text + +if TYPE_CHECKING: + from datetime import datetime + + import asyncpg + + +# The exact set ``ck_jobs_status`` allows. Kept as one expression over the +# shared tuples so the constraint, the migration and this guard cannot drift. +_ALLOWED_STATUSES: frozenset[str] = frozenset(ACTIVE_JOB_STATES) | frozenset(TERMINAL_JOB_STATES) + +_COLUMNS = ( + "id", + "status", + "partition", + "file_id", + "user_id", + "job_metadata", + "error", + "created_at", + "started_at", + "completed_at", + "updated_at", +) + +# Columns ``update_job(**fields)`` may write. ``id`` and ``created_at`` are +# immutable identity; ``updated_at`` is stamped by the repository itself. +_UPDATABLE_COLUMNS = frozenset( + { + "status", + "partition", + "file_id", + "user_id", + "job_metadata", + "error", + "started_at", + "completed_at", + } +) +_SELECT = f"SELECT {', '.join(_COLUMNS)} FROM jobs" -class PgJobRepository(_StubRepositoryBase, JobRepository): - """TODO: real impl once the ``jobs`` table is added. See REFACTORING P0 plan.""" + +class PgJobRepository(JobRepository): + """Store indexation job lifecycle state in Postgres.""" + + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + # ------------------------------------------------------------------ + # Writes + # ------------------------------------------------------------------ async def create_job(self, job: IndexationJob) -> IndexationJob: - raise stub_not_implemented("DB-backed job tracking") + row = await self.pool.fetchrow( + f""" + INSERT INTO jobs (id, status, partition, file_id, user_id, job_metadata, error, + created_at, started_at, completed_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) + ON CONFLICT (id) DO UPDATE + SET status = EXCLUDED.status, + updated_at = now() + RETURNING {", ".join(_COLUMNS)} + """, + job.id, + _status_value(job.status), + job.partition, + job.file_id, + job.user_id, + dict(job.job_metadata), + truncate_error_text(job.error), + job.created_at, + job.started_at, + job.completed_at, + job.updated_at, + ) + created = _row_to_job(row) + if created is None: # pragma: no cover - RETURNING always yields the upserted row + raise RuntimeError(f"jobs upsert returned no row for job_id={job.id!r}") + return created + + async def update_job(self, job_id: str, **fields: Any) -> IndexationJob | None: + """Patch a job row; unknown keys are ignored rather than rejected. + + Callers live on the indexing hot path and pass whatever they know about + the transition, so an unexpected key must not raise there — the + allowlist silently drops it. With nothing left to write this degrades to + a plain read, which keeps ``update_job`` total (it always returns the + current row, or ``None`` if the job is gone). + """ + updates = {key: value for key, value in fields.items() if key in _UPDATABLE_COLUMNS} + if "error" in updates: + updates["error"] = truncate_error_text(updates["error"]) + if "status" in updates: + updates["status"] = _status_value(updates["status"]) + if "job_metadata" in updates: + updates["job_metadata"] = dict(updates["job_metadata"] or {}) + + if not updates: + return await self.get_job(job_id) + + assignments = [f"{column} = ${i}" for i, column in enumerate(updates, start=2)] + assignments.append("updated_at = now()") + + # A cancellation is sticky here for the same reason it is sticky in the + # hot cache (``TaskStateManager.set_state``): once the user has asked + # for a cancel and been told it landed, no later write may walk it back. + # Without this guard a *winning* cancel — one that claims a task still + # SERIALIZING — is clobbered by the worker's own in-flight lifecycle + # write, because both are blind UPDATEs and Postgres is free to order + # them either way. The COMPLETED case leaves the actor and the table + # disagreeing forever; the SERIALIZING case is worse, stranding the row + # non-terminal after ``ray.cancel`` has already killed the only writer + # that could have finished it — and ``purge_terminal_jobs`` sweeps + # terminal rows only, so it never ages out. + # ``mark_failed_if_not_cancelled`` already arbitrates this in SQL; this + # extends the same rule to the writes that lacked it. + # ``None`` covers the status-less patch (a late error or timestamp), + # which deliberately keeps the guard *off*: the rule is about status + # transitions, not about freezing the row — see + # ``test_a_non_status_patch_still_lands_on_a_cancelled_job``. Since + # ``_status_value`` now rejects any value outside the CHECK, an explicit + # ``status=None`` can no longer reach this line, so the two cases the + # ``get`` conflates are down to the one that is intended. + guard_cancelled = updates.get("status") not in (None, "CANCELLED") + where = "id = $1" + (" AND status <> 'CANCELLED'" if guard_cancelled else "") + + row = await self.pool.fetchrow( + f""" + UPDATE jobs + SET {", ".join(assignments)} + WHERE {where} + RETURNING {", ".join(_COLUMNS)} + """, + job_id, + *updates.values(), + ) + if row is None and guard_cancelled: + # The guard declined the write, or the job is gone. Re-read so the + # caller still gets the truth: returning ``None`` here would report + # a cancelled job as a missing one. + return await self.get_job(job_id) + return _row_to_job(row) + + async def mark_failed_if_not_cancelled(self, job_id: str, *, error: str, completed_at: datetime) -> bool: + """Move the row to FAILED unless a cancellation already claimed it. + + The guard is part of the UPDATE, so a concurrent cancel either lands + before this statement (and the WHERE drops it) or after (and it + overwrites a FAILED row the user never saw). No row returned means the + job was already CANCELLED, or is gone. + """ + row = await self.pool.fetchrow( + """ + UPDATE jobs + SET status = 'FAILED', + error = $2, + completed_at = $3, + updated_at = now() + WHERE id = $1 + AND status <> 'CANCELLED' + RETURNING id + """, + job_id, + truncate_error_text(error), + completed_at, + ) + return row is not None + + async def purge_terminal_jobs(self, *, older_than_seconds: int, keep_last: int) -> int: + if older_than_seconds < 0 or keep_last < 0: + raise ValueError("older_than_seconds and keep_last must be non-negative") + + # One statement for both bounds: age evicts the long tail on a busy + # deployment, ``keep_last`` caps the table on a deployment that indexes + # faster than the age window retires rows. ``completed_at IS NULL`` can + # only happen for a row whose terminal write raced a schema/write error, + # so age it out on created_at rather than leaking it forever. + purged = await self.pool.fetchval( + """ + WITH terminal AS ( + SELECT id, + COALESCE(completed_at, created_at) AS settled_at, + row_number() OVER (ORDER BY COALESCE(completed_at, created_at) DESC) AS recency + FROM jobs + WHERE status = ANY($1::text[]) + ), + deleted AS ( + DELETE FROM jobs + WHERE id IN ( + SELECT id FROM terminal + WHERE settled_at < now() - make_interval(secs => $2::double precision) + OR recency > $3 + ) + RETURNING 1 + ) + SELECT COUNT(*)::int FROM deleted + """, + list(TERMINAL_JOB_STATES), + older_than_seconds, + keep_last, + ) + return purged or 0 + + # ------------------------------------------------------------------ + # Reads + # ------------------------------------------------------------------ async def get_job(self, job_id: str) -> IndexationJob | None: - raise stub_not_implemented("DB-backed job tracking") + row = await self.pool.fetchrow(f"{_SELECT} WHERE id = $1", job_id) + return _row_to_job(row) async def list_jobs( self, status: str | None = None, offset: int = 0, limit: int = 50, + user_id: int | None = None, ) -> list[IndexationJob]: - raise stub_not_implemented("DB-backed job tracking") + where = [] + params: list[Any] = [] + states = _expand_status(status) + if states is not None: + params.append(states) + where.append(f"status = ANY(${len(params)}::text[])") + if user_id is not None: + params.append(user_id) + where.append(f"user_id = ${len(params)}") - async def update_job(self, job_id: str, **fields: Any) -> IndexationJob | None: - raise stub_not_implemented("DB-backed job tracking") + clause = f" WHERE {' AND '.join(where)}" if where else "" + params.append(max(0, limit)) + limit_param = len(params) + params.append(max(0, offset)) + + rows = await self.pool.fetch( + f"{_SELECT}{clause} ORDER BY created_at DESC, id DESC LIMIT ${limit_param} OFFSET ${len(params)}", + *params, + ) + return [_row_to_job(row) for row in rows] + + async def count_by_status(self) -> dict[str, int]: + rows = await self.pool.fetch("SELECT status, COUNT(*)::int AS count FROM jobs GROUP BY status") + return {row["status"]: row["count"] for row in rows} + + +def _expand_status(status: str | None) -> list[str] | None: + """Resolve a status filter to the concrete states it selects.""" + if status is None: + return None + if status.lower() == "active": + return list(ACTIVE_JOB_STATES) + return [status.upper()] + + +def _status_value(status: Any) -> str: + """Normalize a status to the exact spelling the ``ck_jobs_status`` CHECK allows. + + Upper-cased for the same reason :func:`_expand_status` does it: a lower-case + string would fail the CHECK, and every durable write is best-effort, so the + violation would be swallowed and the job would silently freeze at its + previous status — permanently, if the dropped write was the terminal one. + Enum members already carry the right value; plain strings may not. + + Membership is checked for the same reason, one step earlier. Casing is not + the only way to build a status the CHECK rejects: ``update_job`` is untyped + and takes whatever the hot path knows, so ``None`` (``str(None).upper()`` is + ``"NONE"``) or a stray :class:`JobStatus` member — the enum this field used + before #660, still exported from ``catalog`` — reaches SQL and violates the + constraint just the same, with the same silent freeze. Raising here turns an + unwritable value into a loud, local failure instead of a lost write. + """ + value = status.value if hasattr(status, "value") else str(status) + value = value.upper() + if value not in _ALLOWED_STATUSES: + raise ValueError( + f"unknown job status {status!r}: the ck_jobs_status CHECK allows only {sorted(_ALLOWED_STATUSES)}" + ) + return value + + +def _row_to_job(row: asyncpg.Record | None) -> IndexationJob | None: + if row is None: + return None + data = dict(row) + data["job_metadata"] = data.get("job_metadata") or {} + return IndexationJob(**data) __all__ = ["PgJobRepository"] diff --git a/openrag/services/persistence/migrations/alembic/versions/d4e5f6a7b8c9_add_jobs_table.py b/openrag/services/persistence/migrations/alembic/versions/d4e5f6a7b8c9_add_jobs_table.py new file mode 100644 index 000000000..d899e7d07 --- /dev/null +++ b/openrag/services/persistence/migrations/alembic/versions/d4e5f6a7b8c9_add_jobs_table.py @@ -0,0 +1,81 @@ +"""add jobs table for durable indexation job state + +Revision ID: d4e5f6a7b8c9 +Revises: b7c1d2e3f4a5 +Create Date: 2026-07-16 10:12:44.118203 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from schema_helpers import index_exists, table_exists +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision: str = "d4e5f6a7b8c9" +down_revision: str | Sequence[str] | None = "b7c1d2e3f4a5" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +_INDEXES = ( + ("ix_jobs_status_created_at", ["status", "created_at"]), + ("ix_jobs_user_status", ["user_id", "status"]), + # Expression index: the retention sweep filters and orders on + # ``COALESCE(completed_at, created_at)``, not on ``completed_at`` alone. + ("ix_jobs_settled_at", [sa.text("COALESCE(completed_at, created_at)")]), +) + + +def upgrade() -> None: + """Upgrade schema. + + Idempotent: ``Base.metadata.create_all()`` runs at app startup, so a freshly + bootstrapped database already has ``jobs`` before alembic reaches this + revision — an unguarded CREATE TABLE would raise ``DuplicateTable``. + """ + if not table_exists("jobs"): + op.create_table( + "jobs", + sa.Column("id", sa.String(), nullable=False), + sa.Column("status", sa.String(), nullable=False), + sa.Column("partition", sa.String(), nullable=False), + sa.Column("file_id", sa.String(), nullable=True), + sa.Column("user_id", sa.Integer(), nullable=True), + sa.Column( + "job_metadata", + postgresql.JSONB(astext_type=sa.Text()), + server_default=sa.text("'{}'::jsonb"), + nullable=False, + ), + sa.Column("error", sa.String(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("started_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.CheckConstraint( + "status IN ('QUEUED','SERIALIZING','CHUNKING','INSERTING','COMPLETED','FAILED','CANCELLED')", + name="ck_jobs_status", + ), + # No FK to partitions.partition: a job row is a historical record and + # must outlive the partition it targeted. + # Constraints stay unnamed so this CREATE TABLE and the startup + # ``metadata.create_all()`` converge on the same Postgres-default + # names (``jobs_pkey`` / ``jobs_user_id_fkey``). + sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="SET NULL"), + sa.PrimaryKeyConstraint("id"), + ) + for name, columns in _INDEXES: + if not index_exists("jobs", name): + op.create_index(name, "jobs", columns) + + +def downgrade() -> None: + """Downgrade schema.""" + for name, _ in _INDEXES: + if index_exists("jobs", name): + op.drop_index(name, table_name="jobs") + if table_exists("jobs"): + op.drop_table("jobs") diff --git a/openrag/services/persistence/schema.py b/openrag/services/persistence/schema.py index f17af7d8a..b824c00d2 100644 --- a/openrag/services/persistence/schema.py +++ b/openrag/services/persistence/schema.py @@ -176,6 +176,65 @@ ) +# Durable indexation-job records (issue #660). Previously job state lived only +# in the detached ``TaskStateManager`` Ray actor, so an API restart made every +# in-flight batch unobservable and un-cancellable. One row per dispatched task; +# ``id`` is the dispatcher's ``task_id``. +# +# ``partition`` deliberately carries no FK to ``partitions.partition``: a job +# row is a historical audit record and must outlive the partition it targeted +# (and a terminal job must not block a partition delete). ``user_id`` mirrors +# ``files.created_by`` with ``ondelete="SET NULL"`` for the same reason. +# +# Rows are bounded by retention, not by the table (see +# ``PgJobRepository.purge_terminal_jobs``). +jobs = Table( + "jobs", + metadata, + Column("id", String, primary_key=True), + Column("status", String, nullable=False), + Column("partition", String, nullable=False), + Column("file_id", String, nullable=True), + Column( + "user_id", + Integer, + ForeignKey("users.id", ondelete="SET NULL"), + nullable=True, + ), + Column("job_metadata", JSONB, server_default=text("'{}'::jsonb"), nullable=False), + Column("error", String, nullable=True), + Column( + "created_at", + DateTime(timezone=True), + server_default=text("now()"), + nullable=False, + ), + Column("started_at", DateTime(timezone=True), nullable=True), + Column("completed_at", DateTime(timezone=True), nullable=True), + Column( + "updated_at", + DateTime(timezone=True), + server_default=text("now()"), + nullable=False, + ), + CheckConstraint( + "status IN ('QUEUED','SERIALIZING','CHUNKING','INSERTING','COMPLETED','FAILED','CANCELLED')", + name="ck_jobs_status", + ), + # Admin queue views filter by status and order by recency. + Index("ix_jobs_status_created_at", "status", "created_at"), + # Per-user task listing and the pending-task count. + Index("ix_jobs_user_status", "user_id", "status"), + # Retention sweeps terminal rows by *settle* time, which is + # ``COALESCE(completed_at, created_at)`` — a row whose terminal write raced a + # failure has no ``completed_at`` and is aged out on ``created_at`` instead. + # The index has to match that expression: a plain b-tree on bare + # ``completed_at`` can serve neither the filter nor the ordering the sweep + # uses, so it would be maintained on every insert and read by nothing. + Index("ix_jobs_settled_at", text("COALESCE(completed_at, created_at)")), +) + + users = Table( "users", metadata, @@ -304,6 +363,7 @@ "topic_tags", "partitions", "files", + "jobs", "users", "oidc_sessions", "partition_memberships", diff --git a/openrag/services/persistence/user_repo.py b/openrag/services/persistence/user_repo.py index 88206850b..1caf7a1b5 100644 --- a/openrag/services/persistence/user_repo.py +++ b/openrag/services/persistence/user_repo.py @@ -211,6 +211,59 @@ async def delete_api_key(self, key_id: str) -> bool: "api_keys table is on the post-refactoring roadmap; use users.token until then.", ) + # ── File-quota reserve / release (issue #664) ──────────────────── + + async def try_reserve_file_slot(self, user_id: int, *, default_quota: int) -> int | None: + """Atomically claim one file slot; ``None`` when the quota is full. + + A single conditional UPDATE is the whole point: the read (is there + room?) and the write (take the room) happen in one statement, under + the row lock, so concurrent admits serialise instead of all + observing the same pre-increment ``file_count`` (issue #664). + + The predicate reproduces the resolved-quota semantics documented on + ``AuthService.validate_file_quota``: + + * ``is_admin`` → always granted (admins bypass quotas); + * ``COALESCE(file_quota, $2)`` → a ``NULL`` per-user quota falls back + to the global default, so an *explicit* per-user limit is honored + even when the global default is negative; + * resolved quota ``< 0`` → unlimited; + * otherwise room exists only while ``file_count < resolved quota``. + + Admins and unlimited users are incremented too (unconditionally), so + ``file_count`` stays the truthful per-uploader total it has always + been — release is then symmetric for every caller. + """ + new_count = await self.pool.fetchval( + """ + UPDATE users + SET file_count = file_count + 1 + WHERE id = $1 + AND ( + is_admin + OR COALESCE(file_quota, $2::int) < 0 + OR file_count < COALESCE(file_quota, $2::int) + ) + RETURNING file_count + """, + user_id, + default_quota, + ) + return new_count + + async def release_file_slot(self, user_id: int) -> None: + """Return one reserved slot to the user's budget. + + Clamped at zero so a double release (or a release racing the legacy + delete-path decrement) can never drive ``file_count`` negative and + hand out free quota. + """ + await self.pool.execute( + "UPDATE users SET file_count = GREATEST(file_count - 1, 0) WHERE id = $1", + user_id, + ) + # ── Legacy method names used by the Phase 7C shim ──────────────── async def create_legacy_user( diff --git a/openrag/services/storage/postgres_store.py b/openrag/services/storage/postgres_store.py index e6b0b3f79..58d0c87f9 100644 --- a/openrag/services/storage/postgres_store.py +++ b/openrag/services/storage/postgres_store.py @@ -84,9 +84,10 @@ def __init__(self, config: RDBConfig, *, run_migrations: bool = True) -> None: self._oidc_session_repo = PgOIDCSessionRepository(pool_getter) self._workspace_repo = PgWorkspaceRepository(pool_getter) + self._job_repo = PgJobRepository(pool_getter) + # Stubs — every method raises StubRepositoryError until the matching # table exists. Listed in the post-refactoring roadmap. - self._job_repo = PgJobRepository(pool_getter) self._chunk_repo = PgChunkRepository(pool_getter) self._prompt_repo = PgPromptRepository(pool_getter) self._conversation_repo = PgConversationRepository(pool_getter) diff --git a/openrag/services/workers/dispatcher.py b/openrag/services/workers/dispatcher.py index 678111ecc..790739ea4 100644 --- a/openrag/services/workers/dispatcher.py +++ b/openrag/services/workers/dispatcher.py @@ -1,10 +1,14 @@ from __future__ import annotations +import asyncio +import time import traceback import uuid +from datetime import UTC, datetime from typing import Any from core.indexing.dispatcher import IndexingDispatcher +from core.models.catalog import DocumentStatus, IndexationJob from core.utils.conts import is_internal_metadata_key, strip_internal_metadata from core.utils.logging import get_logger from ray.exceptions import TaskCancelledError @@ -17,6 +21,16 @@ DEFAULT_TIMEOUT = 60.0 _REQUIRE_EXISTING_PARTITION_KWARG = "require_existing_partition" +# Retention for the durable ``jobs`` table (issue #660). Terminal jobs are swept +# opportunistically from the dispatch path rather than by a background task: a +# sweep is only ever needed *because* jobs are being created, and piggybacking on +# dispatch keeps this out of the app lifecycle (no extra task to own, cancel and +# reason about across API replicas). The interval throttle means a burst of a +# thousand uploads still costs one DELETE. +JOB_RETENTION_SECONDS = 7 * 24 * 3600 +JOB_RETENTION_MAX_ROWS = 10_000 +JOB_PURGE_INTERVAL_SECONDS = 300.0 + class WorkerDispatcher(IndexingDispatcher): """Dispatcher that routes new indexing jobs through ``IndexerPool``. @@ -48,6 +62,7 @@ def __init__( workspace_repo: Any, collection: str, timeout: float = DEFAULT_TIMEOUT, + job_repo: Any = None, ) -> None: self._pool = pool self._tsm = task_state_manager @@ -56,6 +71,12 @@ def __init__( self._workspace_repo = workspace_repo self._collection = collection self._timeout = timeout + # Optional so the dispatcher still runs against a catalog store without a + # job repository (and so tests can build one without Postgres). When it + # is absent, job state degrades to the in-memory actor — the pre-#660 + # behaviour. + self._job_repo = job_repo + self._last_job_purge_at: float | None = None async def _call(self, future: Any, task_description: str) -> Any: from services.workers.ray_utils import call_ray_actor_with_timeout @@ -136,6 +157,7 @@ async def dispatch_indexing( embedder_name: str | None = None, require_existing_partition: bool = False, allow_legacy_require_existing_partition_retry: bool = False, + quota_reserved: bool = False, ) -> str: task_id = uuid.uuid4().hex @@ -153,6 +175,33 @@ async def dispatch_indexing( f"in partition {partition!r} is being deleted" ) + # The durable row is written *before* the task is submitted, so a crash + # between submit and the worker's first state write still leaves the job + # visible (as QUEUED) rather than silently in-flight and unobservable. + # + # It is written *after* the admission gate above, not before, because + # #671 made admission refusable: ``_set_queued_details`` returns False + # when a delete fence covers this file, and forces the in-memory state to + # CANCELLED. A job that was never admitted has no work to record, so + # writing QUEUED here would leave a durable row that no worker will ever + # settle -- non-terminal forever, since retention sweeps terminal rows + # only, and counted active in ``/queue/info`` for good. + await self._record_job( + "create", + task_id, + lambda: self._job_repo.create_job( + IndexationJob( + id=task_id, + status=DocumentStatus.QUEUED, + partition=partition, + file_id=metadata.get("file_id"), + user_id=user.get("id") if user else None, + job_metadata=user_metadata, + ) + ), + ) + await self._maybe_purge_jobs() + task: Any | None = None try: submit_kwargs: dict[str, Any] = { @@ -165,6 +214,9 @@ async def dispatch_indexing( "replace": replace, "indexation_config": indexation_config, "embedder_name": embedder_name, + # #664: tells the worker it owns the reserved file slot and must + # release it if the file never reaches the catalog. + "quota_reserved": quota_reserved, } if require_existing_partition: submit_kwargs[_REQUIRE_EXISTING_PARTITION_KWARG] = True @@ -174,6 +226,20 @@ async def dispatch_indexing( allow_legacy_retry=allow_legacy_require_existing_partition_retry, ) + # #671 made this call fail closed: it returns False when a delete + # fence covers the file, and the handler below then cancels the + # worker we just started and sweeps its vectors. + # + # This is where 4ec0a634 ("don't report a dispatch failure once the + # worker has started") used to swallow the failure, on the grounds + # that the worker owns the reserved slot from submit onwards and + # would run to completion regardless. #671 invalidated that premise: + # the worker no longer runs to completion, it is rolled back, so the + # error is truthful and must propagate. The residual cost is a + # double release (the cancelled worker's finally and the request + # teardown both give the slot back), which under-counts rather than + # leaks -- the direction this branch already chose deliberately, and + # which #700 closes. registered = await self._call( self._tsm.set_object_ref.remote(task_id, {"ref": task}), task_description=f"set_object_ref({task_id})", @@ -238,6 +304,16 @@ async def _submit_indexing_task_once(self, task_id: str, submit_kwargs: dict[str # (wrapped so Ray doesn't auto-dereference and block on the worker task). # Awaiting the submit call yields that list; element 0 is the worker ref # that ``cancel_task``/``ray.cancel`` must target. + # + # A timeout here is ambiguous about ownership of the reserved slot: the + # submit may have started the worker (which then owns it) or not (in + # which case the request's teardown must release it), and the caller + # cannot tell which. The wrapper cancels the future and raises, so the + # router skips ``commit_quota_reservation`` and teardown releases -- + # correct for the overwhelmingly common case that submit never ran, and + # off by one the other way. Both directions undercount rather than + # leak, which is the side this branch errs on, and both self-heal under + # the #676 recount. submitted = await self._call( self._pool.submit.remote(**submit_kwargs), task_description=f"submit({task_id})", @@ -245,16 +321,53 @@ async def _submit_indexing_task_once(self, task_id: str, submit_kwargs: dict[str return submitted[0] async def _mark_submit_failed(self, task_id: str, tb: str) -> None: - set_failed = getattr(self._tsm, "set_failed_if_not_cancelled", None) - if set_failed is not None: - await self._call( - set_failed.remote(task_id, tb), - task_description=f"set_failed_if_not_cancelled({task_id})", + """Settle a dispatch that failed, in both stores. + + The durable half is #660's: ``dispatch_indexing`` has already written a + QUEUED row, and a failed submit means ``IndexerWorker`` -- the only + writer of SERIALIZING/COMPLETED/FAILED -- is never entered, so nothing + else would ever settle it. Left QUEUED the row is unsweepable + (``purge_terminal_jobs`` takes terminal rows only) and counted active in + ``/queue/info`` for good, and the actor entry is unevictable (eviction + reads ``terminal_at``, which only a terminal state enters) on an actor + that outlives the API restart. + + Both writes inherit this method's caller-side gate: #671 only calls it + once ``_cancel_submitted_task`` has confirmed the task really was + cancelled, so a task that actually completed is never recorded FAILED in + either store. + """ + # Guarded: this runs inside the except in dispatch_indexing, so + # an unreachable state actor here would replace the exception the caller + # actually needs ("the pool is down") with one about the bookkeeping + # ("the actor is gone"). The durable write below is guarded by + # _record_job for the same reason. + try: + set_failed = getattr(self._tsm, "set_failed_if_not_cancelled", None) + if set_failed is not None: + await self._call( + set_failed.remote(task_id, tb), + task_description=f"set_failed_if_not_cancelled({task_id})", + ) + else: + await self._call( + self._tsm.set_state.remote(task_id, "FAILED"), + task_description=f"set_state({task_id}, FAILED)", + ) + except Exception as exc: # noqa: BLE001 - settling must not mask the dispatch failure + logger.warning( + "Could not settle a failed dispatch in the state actor; the task may stay QUEUED", + task_id=task_id, + error=str(exc), ) - return - await self._call( - self._tsm.set_state.remote(task_id, "FAILED"), - task_description=f"set_state({task_id}, FAILED)", + await self._record_job( + "settle_failed_dispatch", + task_id, + lambda: self._job_repo.mark_failed_if_not_cancelled( + task_id, + error=tb, + completed_at=datetime.now(UTC), + ), ) async def _cancel_submitted_task(self, task_id: str, task: Any) -> bool: @@ -298,6 +411,52 @@ async def _cancel_submitted_task(self, task_id: str, task: Any) -> bool: ) return False + async def _record_job(self, action: str, task_id: str, call: Any) -> Any: + """Run a durable job write, degrading to a warning on failure. + + Postgres is the source of truth for job state, but it is not on the + critical path of *indexing*: failing an upload because the audit row + could not be written would turn a monitoring outage into a data-ingest + outage. The in-memory actor still has the state, so we log loudly and + continue. + """ + if self._job_repo is None: + return None + try: + return await call() + except Exception as exc: # noqa: BLE001 - durable bookkeeping must not fail indexing + logger.warning( + "Durable job state write failed; job history for this task may be incomplete", + action=action, + task_id=task_id, + error=str(exc), + ) + return None + + async def _maybe_purge_jobs(self) -> None: + """Sweep terminal jobs at most once per ``JOB_PURGE_INTERVAL_SECONDS``.""" + if self._job_repo is None: + return + now = time.monotonic() + if self._last_job_purge_at is not None and now - self._last_job_purge_at < JOB_PURGE_INTERVAL_SECONDS: + return + # Stamped before the call so a slow or failing purge cannot be retried on + # every single dispatch. + self._last_job_purge_at = now + purged = await self._record_job( + "purge", + "-", + lambda: self._job_repo.purge_terminal_jobs( + older_than_seconds=JOB_RETENTION_SECONDS, + keep_last=JOB_RETENTION_MAX_ROWS, + ), + ) + if purged: + logger.info("Purged terminal indexation jobs past retention", purged=purged) + + async def _get_job(self, task_id: str) -> Any: + return await self._record_job("get", task_id, lambda: self._job_repo.get_job(task_id)) + async def delete_file(self, file_id: str, partition: str) -> None: await self._begin_file_delete_fence(file_id=file_id, partition=partition) delete_failed = False @@ -382,14 +541,15 @@ async def copy_file( metadata: dict, partition: str, user: dict | None, - ) -> None: + ) -> bool: + """Copy the file's chunks + catalog row; return whether a row was created.""" rows = await self._vector_store.query_chunks_by_filter( self._collection, {"partition": partition, "file_id": file_id}, output_fields=["*", "vector"], ) if not rows: - return + return False public_metadata = strip_internal_metadata(metadata) entities = [] @@ -405,13 +565,15 @@ async def copy_file( target_partition = metadata.get("partition", partition) file_metadata = self._file_metadata_from_chunk(rows[0]) file_metadata.update(public_metadata) - await self._document_repo.add_file_to_partition( - file_id=target_file_id, - partition=target_partition, - file_metadata=file_metadata, - user_id=user.get("id") if user else None, - relationship_id=file_metadata.get("relationship_id"), - parent_id=file_metadata.get("parent_id"), + return bool( + await self._document_repo.add_file_to_partition( + file_id=target_file_id, + partition=target_partition, + file_metadata=file_metadata, + user_id=user.get("id") if user else None, + relationship_id=file_metadata.get("relationship_id"), + parent_id=file_metadata.get("parent_id"), + ) ) async def _upsert_entities(self, entities: list[dict[str, Any]]) -> None: @@ -434,16 +596,30 @@ def _file_metadata_from_chunk(self, chunk: dict[str, Any]) -> dict[str, Any]: } async def get_task_state(self, task_id: str) -> str | None: - return await self._call( + """Read a task's state, hot cache first, Postgres second. + + A miss is not "unknown task": the actor evicts settled tasks and loses + everything on restart, so the durable row is what makes a task's outcome + observable afterwards. + """ + state = await self._call( self._tsm.get_state.remote(task_id), task_description=f"get_state({task_id})", ) + if state is not None: + return state + job = await self._get_job(task_id) + return job.status.value if job else None async def get_task_error(self, task_id: str) -> str | None: - return await self._call( + error = await self._call( self._tsm.get_error.remote(task_id), task_description=f"get_error({task_id})", ) + if error is not None: + return error + job = await self._get_job(task_id) + return job.error if job else None async def cancel_task(self, task_id: str) -> bool: import ray @@ -465,9 +641,57 @@ async def cancel_task(self, task_id: str) -> bool: task_description=f"set_cancelled_if_active({task_id})", ) if not cancelled: + # Already terminal (or evicted): the task reached its own outcome + # first. Returning before the durable write is what keeps a job that + # actually COMPLETED from being recorded as CANCELLED in `jobs`. return False - ray.cancel(obj_ref["ref"], recursive=True) + # The durable write is part of the claim, so it happens *before* + # ``ray.cancel`` -- not after. ``ray.cancel`` kills the only other + # writer of this row, and everything between here and there runs + # without a successor that could heal a half-applied cancel: + # ``_record_job`` catches ``Exception``, but a client disconnect or a + # graceful shutdown raises ``asyncio.CancelledError``, a + # ``BaseException`` that sails straight through it. Writing after the + # kill would leave the actor CANCELLED and the row stuck on its last + # active status forever -- non-terminal, so ``purge_terminal_jobs`` + # never sweeps it, ``count_by_status`` counts it active forever, and + # the actor-first and durable-first read paths answer differently for + # the same task id. + # + # Ordering it first cannot mis-record a job that escapes the cancel: + # the actor claim above already succeeded, and ``update_job`` keeps + # CANCELLED sticky, so a worker that somehow finishes anyway is + # declined by the same guard in both stores. + try: + # ``shield`` so a client disconnect cannot abort the UPDATE in + # flight: the write runs to completion even though the + # ``CancelledError`` propagates to us immediately. + await asyncio.shield( + self._record_job( + "cancel", + task_id, + lambda: self._job_repo.update_job( + task_id, + status=DocumentStatus.CANCELLED, + completed_at=datetime.now(UTC), + ), + ) + ) + finally: + # In a ``finally`` so the worker still dies if the durable write + # raises: the actor has already claimed the cancellation, and + # leaving the worker running would contradict both records. + ray.cancel(obj_ref["ref"], recursive=True) + # The reserved quota slot (#664) is deliberately not released here. + # A task cancelled mid-flight gives its slot back in + # ``IndexerWorker.process_file``'s ``finally``; releasing here too + # would double-release. But a task that ``ray.cancel`` retires + # *before* that body runs never executes the ``finally``, so its slot + # leaks -- and the CANCELLED row written just above is terminal, hence + # indistinguishable from a clean cancel. Recovering it needs the #676 + # reconciliation to *recount* ``file_count`` (completed files + active + # job rows), not merely sweep orphaned active rows. return True @@ -479,6 +703,7 @@ def from_ray_namespace( document_repo: Any, workspace_repo: Any, collection: str, + job_repo: Any = None, ) -> WorkerDispatcher: import ray from services.workers.indexer_pool import build_indexer_pool @@ -491,6 +716,7 @@ def from_ray_namespace( workspace_repo=workspace_repo, collection=collection, timeout=timeout, + job_repo=job_repo, ) diff --git a/openrag/services/workers/indexer_actor.py b/openrag/services/workers/indexer_actor.py index dd84b12e8..6bb75507f 100644 --- a/openrag/services/workers/indexer_actor.py +++ b/openrag/services/workers/indexer_actor.py @@ -2,10 +2,11 @@ import asyncio import traceback -from datetime import datetime +from datetime import UTC, datetime from pathlib import Path from typing import Any +from core.models.catalog import DocumentStatus from core.models.document import Document from core.utils.logging import get_logger from services.workers.pipeline_builder import ( @@ -18,6 +19,8 @@ logger = get_logger() +logger = get_logger() + class IndexerWorker: """Pure-Python core of the thin indexer actor. @@ -35,6 +38,14 @@ class IndexerWorker: Callers are responsible for setting ``QUEUED`` *before* dispatching the task, and for storing the object ref via ``set_object_ref``. + + Since #664 the worker is also the owner of the uploader's reserved + quota slot: when the caller admitted the upload it already charged one + ``users.file_count`` slot, and this worker either consumes it (by + writing the catalog row) or releases it — see ``process_file``. + + Every transition is mirrored to *job_repo* (issue #660) so the outcome of a + file survives a restart of this worker or of the state actor. """ def __init__( @@ -45,6 +56,8 @@ def __init__( topic_tag_repo: Any = None, vector_store: Any = None, collection: str = "default", + user_repo: Any = None, + job_repo: Any = None, ) -> None: self._pipeline = pipeline self._tsm = task_state_manager @@ -52,6 +65,8 @@ def __init__( self._topic_tag_repo = topic_tag_repo self._vector_store = vector_store self._collection = collection + self._user_repo = user_repo + self._job_repo = job_repo async def process_file( self, @@ -66,21 +81,45 @@ async def process_file( indexation_config: dict[str, Any] | None = None, embedder_name: str | None = None, require_existing_partition: bool = False, + quota_reserved: bool = False, ) -> dict[str, Any]: """Run one file through the indexing pipeline. Returns a plain dict ``{"stored_count": int, "stage": "stored"}`` on success. On failure, state is set to FAILED and the exception is re-raised so the Ray task is marked as errored. + + When ``quota_reserved`` is set, the admission gate already charged + one ``users.file_count`` slot to the uploader (#664) and this call + owns it. The slot is *consumed* the moment ``add_file_to_partition`` + reports a new catalog row; on every other outcome — pipeline + failure, cancellation, or the duplicate-at-catalog race where the + insert reports ``False`` — the ``finally`` below hands it back, so a + rejected upload can never permanently eat a slot. """ - await self._tsm.set_state.remote(task_id, "SERIALIZING") + # Released in ``finally`` unless the catalog write claims it. A flag + # plus ``finally`` (rather than an ``except``) is what makes + # cancellation safe: ray.cancel raises ``asyncio.CancelledError``, a + # BaseException that ``except Exception`` would sail straight past. + release_slot = bool(quota_reserved) + user_id = (user or {}).get("id") row: dict[str, Any] | None = None catalog_written = False try: + # Inside the release scope: ``set_state`` talks to a detached actor + # that can be unreachable, and by this point the request has already + # handed the slot over at dispatch, so nothing else would give it back. + await self._tsm.set_state.remote(task_id, "SERIALIZING") + await _update_job( + self._job_repo, + task_id, + status=DocumentStatus.SERIALIZING, + started_at=datetime.now(UTC), + ) document = await _load_document(path, metadata, partition) # One indexation timestamp for this file, shared by the Milvus chunks # (via the store stage) and the Postgres catalog row, so they agree. - row: dict[str, Any] = { + row = { "task_id": task_id, "document": document, "partition": partition, @@ -114,6 +153,20 @@ async def process_file( row=row, timeout=getattr(getattr(self._pipeline, "timeouts", None), "store", None), ) + if not replace: + # A *new* catalog row now exists for this reservation, so it + # is consumed; releasing it would hand the user free quota. + # ``replace`` reuses an existing row and creates nothing, but + # it also never reserves (only ``add_file`` does), so there is + # no slot to consume on that branch. + # + # The duplicate-at-catalog race no longer reaches here at all: + # ``add_file_to_partition`` returns False for an existing row, + # which the fail-closed check above turns into a raise, and + # ``process_file``'s ``finally`` releases the slot on the way + # out. Same outcome as the pre-rebase ``if created:`` gate, + # reached by a different path. + release_slot = False if self._topic_tag_repo is not None: await _replace_topic_tags_if_needed( topic_tag_repo=self._topic_tag_repo, @@ -123,6 +176,12 @@ async def process_file( indexation_config=indexation_config, ) await self._tsm.set_state.remote(task_id, "COMPLETED") + await _update_job( + self._job_repo, + task_id, + status=DocumentStatus.COMPLETED, + completed_at=datetime.now(UTC), + ) return {"stored_count": row.get("stored_count", 0), "stage": row.get("stage", "")} except Exception: should_cleanup_vectors = row is not None and ( @@ -137,14 +196,110 @@ async def process_file( task_id=task_id, ) tb = traceback.format_exc() - await self._tsm.set_failed_if_not_cancelled.remote(task_id, tb) + # Keep the hot cache in step, but do not let its verdict decide the + # durable write. ``set_failed_if_not_cancelled`` returns False both + # for a real cancellation *and* for an entry the actor no longer has + # (restart, TTL eviction, lost node), and skipping the write in the + # second case strands the row in SERIALIZING forever — retention + # sweeps terminal rows only. Postgres arbitrates instead: it is the + # one participant guaranteed to still know what the user asked for. + # Which is why this one is guarded: unguarded, an actor that is gone + # (restart, lost node) raises straight out of this handler and takes + # ``_mark_job_failed`` with it, stranding the row exactly as above — + # so the hot cache would decide the durable outcome after all. + try: + await self._tsm.set_failed_if_not_cancelled.remote(task_id, tb) + except Exception as exc: # noqa: BLE001 - hot-cache write must not block the durable one + logger.warning( + "Hot-cache FAILED write lost; the durable job row arbitrates.", + task_id=task_id, + error=str(exc), + ) + await _mark_job_failed( + self._job_repo, + task_id, + error=tb, + completed_at=datetime.now(UTC), + ) raise + finally: + if release_slot: + await release_quota_slot(self._user_repo, user_id) # The raw upload is purged (when configured) by the enclosing actor, not # here: cleanup must also cover failures that happen *before* this method # runs (catalog/registry init, the SERIALIZING state update). See # ``delete_uploaded_file`` and ``IndexerWorkerActor.process_file``. +async def _update_job(job_repo: Any, task_id: str, **fields: Any) -> None: + """Mirror a lifecycle transition to the durable ``jobs`` row. + + Best-effort by design: this is bookkeeping about the work, not the work. A + Postgres blip must not fail a file that indexed correctly (nor mask the real + exception on the failure path, where this runs inside an ``except`` block). + The repository truncates the stored traceback. + """ + if job_repo is None: + return + try: + await job_repo.update_job(task_id, **fields) + except Exception as exc: # noqa: BLE001 - durable bookkeeping must not fail indexing + logger.warning( + "Durable job state write failed; job history for this task may be incomplete", + task_id=task_id, + status=fields.get("status"), + error=str(exc), + ) + + +async def _mark_job_failed(job_repo: Any, task_id: str, *, error: str, completed_at: datetime) -> None: + """Mirror a FAILED outcome to the durable row, with the cancel check in SQL. + + Same best-effort contract as :func:`_update_job` — a Postgres blip must not + mask the exception this runs inside. The CANCELLED guard is part of the + UPDATE, so a state actor that lost the task cannot leave the row non-terminal. + """ + if job_repo is None: + return + try: + await job_repo.mark_failed_if_not_cancelled(task_id, error=error, completed_at=completed_at) + except Exception as exc: # noqa: BLE001 - durable bookkeeping must not fail indexing + logger.warning( + "Durable job state write failed; job history for this task may be incomplete", + task_id=task_id, + status="FAILED", + error=str(exc), + ) + + +async def mark_dispatch_orphan_failed(task_state_manager: Any, job_repo: Any, task_id: str) -> None: + """Drive a task that never reached :class:`IndexerWorker` to a terminal FAILED. + + Used by the pool actor when setup (catalog / model-endpoint registry) blows + up before ``process_file`` is entered. Both stores hold the task as QUEUED + at that point and nothing else will ever write it, so both need settling — + otherwise the entry is unevictable in the detached actor (eviction is driven + off ``terminal_at``, which only a terminal state enters) and the row is + unsweepable by retention, which takes terminal rows only. + + Both writes swallow ``Exception`` so a bookkeeping failure does not mask + the setup exception the caller must actually see; neither swallows + ``BaseException``, so a cancellation arriving here still propagates. + Both carry the CANCELLED guard, so a user who cancelled during setup + keeps their outcome. + """ + error = "Indexing setup failed before the file reached a worker." + try: + await task_state_manager.set_failed_if_not_cancelled.remote(task_id, error) + except Exception as exc: # noqa: BLE001 - settling an orphan must not mask the real failure + logger.warning( + "Hot-cache FAILED write lost while settling a dispatch orphan; the task may stay QUEUED", + task_id=task_id, + error=str(exc), + ) + await _mark_job_failed(job_repo, task_id, error=error, completed_at=datetime.now(UTC)) + + async def _write_catalog_record( *, doc_repo: Any, @@ -156,6 +311,20 @@ async def _write_catalog_record( indexed_at: datetime | None = None, require_existing_partition: bool = False, ) -> bool: + """Write the file's catalog row; return whether the write landed. + + True means the catalog now holds a row for this task -- an updated one on a + ``replace`` re-index, a newly inserted one otherwise. False means the write + did not land, which the caller treats as fatal. + + Note this is *not* "a new file was created". Before the rebase onto #671 + this returned that instead, so a ``replace`` reported False and the caller + read it as "reservation unconsumed". #671 added a fail-closed + ``if not wrote_catalog: raise`` on the same value, which those False returns + would have turned into a hard failure of every ``replace`` re-index. The two + questions are now separate: this one reports the write, and the caller + derives the quota verdict from ``replace`` (#664). + """ file_id = metadata.get("file_id", "") file_metadata = {key: value for key, value in metadata.items() if key != "page"} config_kwargs = {"indexation_config": indexation_config} if indexation_config is not None else {} @@ -312,6 +481,26 @@ def _display_filename(path: str, metadata: dict[str, Any]) -> str: return Path(path).name +async def release_quota_slot(user_repo: Any, user_id: int | None) -> None: + """Give the uploader's reserved file slot back, swallowing any error. + + Runs on cleanup paths only. A release that raises would replace the real + indexing error with a database error, so failures are swallowed — but + loudly, because a lost release leaves ``file_count`` one too high and + permanently narrows that user's quota until an admin reconciles it. + """ + if user_repo is None or user_id is None: + return + try: + await user_repo.release_file_slot(user_id) + except Exception: # noqa: BLE001 - cleanup must never mask the indexing error + logger.exception( + "Failed to release a reserved file slot; the user file_count is now one too high " + "and must be reconciled manually.", + user_id=user_id, + ) + + async def delete_uploaded_file(path: str, logger: Any) -> None: """Remove the raw upload from disk, swallowing any cleanup error. diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index ee89aec82..f75790d9f 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -7,7 +7,12 @@ from typing import Any import ray -from services.workers.indexer_actor import IndexerWorker, delete_uploaded_file +from services.workers.indexer_actor import ( + IndexerWorker, + delete_uploaded_file, + mark_dispatch_orphan_failed, + release_quota_slot, +) from openrag.core.config.root import Settings @@ -62,6 +67,10 @@ def __init__(self) -> None: ) self._vector_store = MilvusVectorStore(cfg.vectordb) task_state_manager = ray.get_actor("TaskStateManager", namespace="openrag") + # Kept on the actor as well as handed to the worker: setup can fail + # before the worker is ever entered, and that path has to settle the + # task itself (see ``process_file``). + self._task_state_manager = task_state_manager pipeline = build_indexing_pipeline( parser=parser, chunker=chunker, @@ -115,6 +124,10 @@ def __init__(self) -> None: topic_tag_repo=self._catalog_store.topic_tag_repo, vector_store=self._vector_store, collection=cfg.vectordb.collection_name, + # #664: the worker owns the uploader's reserved file slot and + # releases it when the file never reaches the catalog. + user_repo=self._catalog_store.user_repo, + job_repo=self._catalog_store.job_repo, ) # When False (e.g. Twake, which keeps its own copy), the raw upload is # purged from ``paths.data_dir`` once indexing settles. Enforced at this @@ -212,10 +225,48 @@ async def process_file( indexation_config: dict[str, Any] | None = None, embedder_name: str | None = None, require_existing_partition: bool = False, + quota_reserved: bool = False, ) -> dict[str, Any]: try: - await self._ensure_catalog() - await self._ensure_registry_fresh(_required_model_endpoint_names(indexation_config, embedder_name)) + try: + await self._ensure_catalog() + await self._ensure_registry_fresh(_required_model_endpoint_names(indexation_config, embedder_name)) + except BaseException: + # Setup blew up before the worker could take ownership of the + # reserved quota slot (#664), so release it here — the worker's + # own finally will never run. ``BaseException`` because a + # cancellation during setup must release too. + # + # Best-effort, and honestly so: this releases through the very + # store whose initialization may have just failed. A cancellation + # or a registry-refresh failure releases fine, but an unreachable + # Postgres — ``_ensure_catalog``'s most likely failure — fails the + # release too, and ``release_quota_slot`` swallows it, leaking the + # slot. Nothing local can fix that (the DB *is* the counter); + # recovering it needs the reconciliation sweep tracked in #676. + if quota_reserved: + await release_quota_slot(self._catalog_store.user_repo, (user or {}).get("id")) + # The task is QUEUED in both stores and ``IndexerWorker`` -- + # the only writer of SERIALIZING/COMPLETED/FAILED -- is never + # entered, so without this it stays non-terminal forever: + # unevictable in the detached actor (eviction reads + # ``terminal_at``, which only a terminal state enters) and + # unsweepable in Postgres (retention takes terminal rows only), + # counted active in ``/queue/info`` for good. The client has + # already been handed a 201 and a ``task_status_url`` that would + # answer QUEUED forever. + # + # Best-effort and, like the release above, honestly so: the + # durable half writes through the very store whose init may have + # just failed. The hot-cache half is independent of it, so an + # unreachable Postgres still leaves an evictable, terminal + # ``TaskInfo`` rather than a pinned one. + await mark_dispatch_orphan_failed( + self._task_state_manager, + self._catalog_store.job_repo, + task_id, + ) + raise result = await self._worker.process_file( task_id=task_id, path=path, @@ -227,6 +278,7 @@ async def process_file( indexation_config=indexation_config, embedder_name=embedder_name, require_existing_partition=require_existing_partition, + quota_reserved=quota_reserved, ) file_id = metadata.get("file_id", "") if workspace_ids and not replace and file_id: diff --git a/openrag/services/workers/task_state.py b/openrag/services/workers/task_state.py index 32716f51a..42c22dfec 100644 --- a/openrag/services/workers/task_state.py +++ b/openrag/services/workers/task_state.py @@ -1,11 +1,14 @@ from __future__ import annotations import asyncio +import time +from collections import OrderedDict from dataclasses import dataclass, field from typing import Any import ray from core.models.catalog import TERMINAL_TASK_STATES, DocumentStatus +from core.utils.text import MAX_ERROR_TEXT_CHARS, truncate_error_text ACTIVE_INDEXING_STATES = frozenset({"QUEUED", "SERIALIZING", "CHUNKING", "INSERTING"}) TERMINAL_INDEXING_STATES = frozenset({"COMPLETED", "FAILED"}) @@ -27,6 +30,32 @@ _MAX_TASKS_PER_WORKER = 1 +# This actor is created with ``lifetime="detached"`` (see ``bootstrap.py``), so it +# outlives the API process and used to be insert-only: every file ever dispatched +# left a permanent ``TaskInfo`` (plus a full traceback on failure) until the actor +# OOM-ed (issue #660). Postgres now holds the durable record, which frees this +# actor to be what its callers actually need — a hot cache of recent tasks. +# +# Only *terminal* tasks are evictable: an in-flight task still owns the +# ``object_ref`` that ``cancel_task`` needs, and that ref is not serializable, so +# it cannot live anywhere but here. In-flight tasks are self-limiting (the pool +# has bounded capacity and every task eventually settles), terminal ones are not. +# +# Both bounds apply, and both are enforced lazily: the sweep runs only when a +# task settles, never on a timer (this is a Ray actor; a background loop would be +# another thing to own). The cap is therefore the real memory guarantee — it is +# checked exactly when growth happens. The TTL only retires stale entries once +# *some* task settles, so a fully idle deployment keeps its last few terminal +# tasks cached indefinitely. That is harmless: a terminal state is immutable, so +# a stale read is not a wrong read, and the cap still bounds the memory. +# Reads that miss fall back to Postgres (``WorkerDispatcher.get_task_state`` / +# ``JobService``), which is the durable record either way. +_TERMINAL_STATES = frozenset({"COMPLETED", "FAILED", "CANCELLED"}) +_MAX_TERMINAL_TASKS = 2000 +_TERMINAL_TTL_SECONDS = 3600.0 +_MAX_ERROR_CHARS = MAX_ERROR_TEXT_CHARS + + @dataclass class TaskInfo: state: str | None = None @@ -41,9 +70,42 @@ def __init__(self) -> None: self.tasks: dict[str, TaskInfo] = {} self.user_index: dict[int | None, set[str]] = {} self.file_delete_fences: dict[tuple[str, str], int] = {} + # task_id -> monotonic timestamp of the terminal transition, in + # insertion order so eviction is FIFO (oldest settled task first). + self.terminal_at: OrderedDict[str, float] = OrderedDict() self.lock = asyncio.Lock() async def _ensure_task(self, task_id: str) -> TaskInfo: + """Create the entry for a task we are hearing about for the first time. + + The call that legitimately means "new" is the dispatcher's opening + ``QUEUED`` write, the first write for a task id (before + ``set_details``/``set_object_ref``). + + ``set_state`` is *not* dispatcher-only, though: the worker also writes + ``SERIALIZING`` and ``COMPLETED`` through it. That is safe today only + because of ordering -- the worker's first write happens long before the + task can be terminal, and eviction only ever removes *terminal* entries, + so there is nothing evicted for it to resurrect. Anything that breaks + that ordering (a ``set_state`` after a terminal transition) reopens the + leak this guard exists to close: a resurrected entry with a + *non-terminal* state never re-enters ``terminal_at`` and is never + evictable again. + + Making creation opt-in (``create=True``, dispatcher only) is the durable + fix, but it changes the signature of a **detached** actor method -- + ``get_or_create_actor(..., lifetime="detached")`` keeps the previous + instance alive across an API deploy, so a new dispatcher would call an + old actor and every dispatch would fail on the unexpected keyword. It + therefore has to be sequenced with a deliberate actor restart rather + than shipped as a plain code change (tracked in #676). + + Every *other* writer must go through :meth:`_live_task`. Creating an + entry from a late write would resurrect an evicted task with + ``state=None``, which never re-enters ``terminal_at`` and is therefore + never evictable again — an unbounded leak on a detached actor, i.e. the + exact failure #660 exists to fix. + """ if task_id not in self.tasks: self.tasks[task_id] = TaskInfo() return self.tasks[task_id] @@ -87,6 +149,55 @@ async def end_file_delete(self, *, partition: str, file_id: str) -> None: else: self.file_delete_fences.pop(key, None) + def _live_task(self, task_id: str) -> TaskInfo | None: + """The entry for ``task_id``, or ``None`` if it is unknown or evicted. + + A write for a task that is no longer cached is dropped: the durable + ``jobs`` row is the record of what happened, and this actor is only a + hot cache of recent tasks. See :meth:`_ensure_task` for why recreating + it here would be a leak. + """ + return self.tasks.get(task_id) + + def _mark_terminal(self, task_id: str, state: str | None) -> None: + """Record (or clear) a task's terminal transition, then evict. + + Called with ``self.lock`` held, from every state write. + """ + if state in _TERMINAL_STATES: + self.terminal_at[task_id] = time.monotonic() + self.terminal_at.move_to_end(task_id) + self._evict_terminal() + else: + # A task that leaves a terminal state (a re-dispatch reusing the id) + # must stop being a candidate for eviction. + self.terminal_at.pop(task_id, None) + + def _evict_terminal(self) -> None: + """Drop terminal tasks that are over the cap or past the TTL.""" + now = time.monotonic() + while self.terminal_at: + task_id, settled_at = next(iter(self.terminal_at.items())) + over_cap = len(self.terminal_at) > _MAX_TERMINAL_TASKS + expired = now - settled_at > _TERMINAL_TTL_SECONDS + if not (over_cap or expired): + # FIFO: the head is the oldest, so nothing behind it can qualify. + break + self.terminal_at.popitem(last=False) + self._forget(task_id) + + def _forget(self, task_id: str) -> None: + info = self.tasks.pop(task_id, None) + if info is None: + return + user_id = info.details.get("user_id") + task_ids = self.user_index.get(user_id) + if task_ids is None: + return + task_ids.discard(task_id) + if not task_ids: + del self.user_index[user_id] + @ray.method(concurrency_group="set") async def set_state(self, task_id: str, state: str) -> None: async with self.lock: @@ -94,12 +205,15 @@ async def set_state(self, task_id: str, state: str) -> None: if info.state == DocumentStatus.CANCELLED and state != DocumentStatus.CANCELLED: return info.state = state + self._mark_terminal(task_id, state) @ray.method(concurrency_group="set") async def set_error(self, task_id: str, tb_str: str) -> None: async with self.lock: - info = await self._ensure_task(task_id) - info.error = tb_str + info = self._live_task(task_id) + if info is None: + return + info.error = truncate_error_text(tb_str, _MAX_ERROR_CHARS) @ray.method(concurrency_group="set") async def set_failed_if_not_cancelled(self, task_id: str, tb_str: str) -> bool: @@ -109,7 +223,8 @@ async def set_failed_if_not_cancelled(self, task_id: str, tb_str: str) -> bool: if info is None or info.state == "CANCELLED": return False info.state = "FAILED" - info.error = tb_str + info.error = truncate_error_text(tb_str, _MAX_ERROR_CHARS) + self._mark_terminal(task_id, "FAILED") return True @ray.method(concurrency_group="set") @@ -119,6 +234,14 @@ async def set_cancelled_if_active(self, task_id: str) -> bool: if info is None or info.state in TERMINAL_TASK_STATES: return False info.state = "CANCELLED" + # CANCELLED is terminal, so it must register like every other + # terminal write does. Skipping this leaves the entry retained + # forever: eviction is driven entirely off ``terminal_at``, and + # nothing writes this task's state again -- ray.cancel raises + # CancelledError, a BaseException that ``process_file``'s + # ``except Exception`` never catches. That is precisely the + # unbounded growth #660 exists to fix. + self._mark_terminal(task_id, "CANCELLED") return True @ray.method(concurrency_group="set") @@ -132,7 +255,11 @@ async def set_details( user_id: int | None, ) -> None: async with self.lock: - info = await self._ensure_task(task_id) + info = self._live_task(task_id) + if info is None: + # Dropping the details also keeps ``user_index`` from growing an + # entry that ``_forget`` will never be able to clean up. + return self._record_details( task_id, info, @@ -173,7 +300,9 @@ async def set_queued_details( @ray.method(concurrency_group="set") async def set_object_ref(self, task_id: str, object_ref: ray.ObjectRef) -> bool: async with self.lock: - info = await self._ensure_task(task_id) + info = self._live_task(task_id) + if info is None: + return info.object_ref = object_ref details = info.details or {} if self._file_delete_fenced(partition=details.get("partition"), file_id=details.get("file_id")): diff --git a/tests/integration/repos/conftest.py b/tests/integration/repos/conftest.py index a0e9dc155..14c98071c 100644 --- a/tests/integration/repos/conftest.py +++ b/tests/integration/repos/conftest.py @@ -1,7 +1,7 @@ """Shared fixtures for the Phase 7F persistence-layer integration tests. The whole suite auto-skips when a Postgres instance is not reachable. We try -the explicit ``POSTGRES_TEST_DSN`` env var first, then fall back to the local +the explicit ``POSTGRES_TEST_ADMIN_DSN`` env var first, then fall back to the local docker-compose ``rdb`` container (the dev DB the rest of the project assumes is up). One ephemeral test database is created at session start, Alembic migrations run once, and individual tests share the same @@ -140,6 +140,7 @@ async def postgres_store(test_rdb_config: RDBConfig) -> PostgresStore: _TRUNCATE_SQL = """ TRUNCATE TABLE + jobs, oidc_sessions, workspace_files, workspaces, diff --git a/tests/integration/repos/test_job_repo.py b/tests/integration/repos/test_job_repo.py new file mode 100644 index 000000000..1ea0f3521 --- /dev/null +++ b/tests/integration/repos/test_job_repo.py @@ -0,0 +1,269 @@ +"""Durable indexation job state against a real Postgres (issue #660). + +Complements the fake-pool unit tests: these are what actually prove the +``jobs`` migration applied, the CHECK constraint matches the state taxonomy, +and the retention sweep deletes what it claims to. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta + +import asyncpg +import pytest +from core.models.catalog import DocumentStatus, IndexationJob +from services.persistence.job_repo import PgJobRepository +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +@pytest.fixture +def repo(postgres_store: PostgresStore) -> PgJobRepository: + return postgres_store.job_repo + + +async def _user(postgres_store: PostgresStore, name: str = "uploader") -> int: + # users.created_at is populated by the ORM/model default, not a server + # default, so a raw INSERT has to supply it. + return await postgres_store.pool.fetchval( + "INSERT INTO users (display_name, is_admin, created_at) VALUES ($1, false, now()) RETURNING id", + name, + ) + + +def _job(**overrides) -> IndexationJob: + data = {"id": "task-1", "partition": "tenant-a", "file_id": "file-1", "job_metadata": {"filename": "a.pdf"}} + data.update(overrides) + return IndexationJob(**data) + + +class TestLifecycle: + async def test_create_then_read_round_trips_every_field(self, repo, postgres_store): + user_id = await _user(postgres_store) + await repo.create_job(_job(user_id=user_id)) + + job = await repo.get_job("task-1") + + assert job.id == "task-1" + assert job.status is DocumentStatus.QUEUED + assert job.partition == "tenant-a" + assert job.file_id == "file-1" + assert job.user_id == user_id + assert job.job_metadata == {"filename": "a.pdf"} + assert job.error is None + + async def test_full_transition_to_completed_is_durable(self, repo): + await repo.create_job(_job()) + started = datetime.now(UTC) + + await repo.update_job("task-1", status=DocumentStatus.SERIALIZING, started_at=started) + await repo.update_job("task-1", status=DocumentStatus.COMPLETED, completed_at=datetime.now(UTC)) + + job = await repo.get_job("task-1") + assert job.status is DocumentStatus.COMPLETED + assert job.started_at is not None + assert job.completed_at is not None + + async def test_failure_stores_a_truncated_traceback(self, repo): + await repo.create_job(_job()) + + await repo.update_job("task-1", status=DocumentStatus.FAILED, error="boom\n" + "x" * 200_000) + + job = await repo.get_job("task-1") + assert job.status is DocumentStatus.FAILED + assert len(job.error) < 10_000 + assert "truncated" in job.error + + async def test_a_late_failure_cannot_overwrite_a_cancellation(self, repo): + await repo.create_job(_job()) + await repo.update_job("task-1", status=DocumentStatus.CANCELLED, completed_at=datetime.now(UTC)) + + wrote = await repo.mark_failed_if_not_cancelled("task-1", error="boom", completed_at=datetime.now(UTC)) + + assert wrote is False + assert (await repo.get_job("task-1")).status is DocumentStatus.CANCELLED + + async def test_a_failure_lands_when_the_job_was_not_cancelled(self, repo): + """The write must not depend on the in-memory actor still knowing the task. + + Gating it on the actor's verdict stranded the row in SERIALIZING whenever + the actor had restarted or evicted the entry, and retention only sweeps + terminal rows — so the job stayed in the queue views forever (#660). + """ + await repo.create_job(_job()) + await repo.update_job("task-1", status=DocumentStatus.SERIALIZING, started_at=datetime.now(UTC)) + + wrote = await repo.mark_failed_if_not_cancelled( + "task-1", error="boom\n" + "x" * 200_000, completed_at=datetime.now(UTC) + ) + + assert wrote is True + job = await repo.get_job("task-1") + assert job.status is DocumentStatus.FAILED + assert job.completed_at is not None + assert len(job.error) < 10_000 # truncation holds on this path too + + async def test_a_late_lifecycle_write_cannot_resurrect_a_cancellation(self, repo): + """A cancel that *wins* must not be walked back by the worker's own write. + + `mark_failed_if_not_cancelled` arbitrated the FAILED path in SQL, but the + worker's SERIALIZING/COMPLETED writes were blind UPDATEs racing the + cancel's blind UPDATE, so Postgres was free to order them either way. + SERIALIZING landing last is the damaging one: `ray.cancel` has already + killed the only writer that could have finished the row, and + `purge_terminal_jobs` sweeps terminal rows only — so the job would sit + in the queue views, active, forever. + """ + await repo.create_job(_job()) + await repo.update_job("task-1", status=DocumentStatus.CANCELLED, completed_at=datetime.now(UTC)) + + declined = await repo.update_job("task-1", status=DocumentStatus.SERIALIZING, started_at=datetime.now(UTC)) + + # The declined write still reports the row, not ``None`` — a cancelled + # job must not read back as a missing one. + assert declined is not None + assert declined.status is DocumentStatus.CANCELLED + assert (await repo.get_job("task-1")).status is DocumentStatus.CANCELLED + + async def test_a_late_completion_cannot_overwrite_a_cancellation(self, repo): + """Otherwise the actor (sticky CANCELLED) and the table disagree forever. + + The two read paths would then answer differently for the same task: + `JobService` reads Postgres first, `WorkerDispatcher.get_task_state` + reads the actor first. + """ + await repo.create_job(_job()) + await repo.update_job("task-1", status=DocumentStatus.CANCELLED, completed_at=datetime.now(UTC)) + + await repo.update_job("task-1", status=DocumentStatus.COMPLETED, completed_at=datetime.now(UTC)) + + assert (await repo.get_job("task-1")).status is DocumentStatus.CANCELLED + + async def test_a_cancellation_is_still_rewritable_as_cancelled(self, repo): + """The guard must not block the cancel path itself (retry, double-click).""" + await repo.create_job(_job()) + first = datetime.now(UTC) + await repo.update_job("task-1", status=DocumentStatus.CANCELLED, completed_at=first) + + later = first + timedelta(seconds=30) + again = await repo.update_job("task-1", status=DocumentStatus.CANCELLED, completed_at=later) + + assert again.status is DocumentStatus.CANCELLED + assert again.completed_at == later + + async def test_a_non_status_patch_still_lands_on_a_cancelled_job(self, repo): + """The guard is about status transitions, not about freezing the row.""" + await repo.create_job(_job()) + await repo.update_job("task-1", status=DocumentStatus.CANCELLED, completed_at=datetime.now(UTC)) + + patched = await repo.update_job("task-1", error="late diagnostic") + + assert patched.status is DocumentStatus.CANCELLED + assert patched.error == "late diagnostic" + + async def test_an_explicit_zero_limit_returns_no_rows(self, repo): + """``limit=0`` means none, not one — the floor is 0, mirroring ``offset``.""" + await repo.create_job(_job()) + + assert await repo.list_jobs(limit=0) == [] + assert len(await repo.list_jobs(limit=1)) == 1 + + async def test_marking_an_unknown_job_failed_reports_no_write(self, repo): + assert await repo.mark_failed_if_not_cancelled("ghost", error="boom", completed_at=datetime.now(UTC)) is False + + async def test_create_is_idempotent_for_a_redispatched_task(self, repo): + await repo.create_job(_job()) + await repo.create_job(_job(status=DocumentStatus.CANCELLED)) + + assert (await repo.get_job("task-1")).status is DocumentStatus.CANCELLED + assert await postgres_count(repo) == 1 + + async def test_update_of_an_unknown_job_returns_none(self, repo): + assert await repo.update_job("ghost", status=DocumentStatus.FAILED) is None + + async def test_an_unknown_state_is_rejected_by_the_check_constraint(self, repo): + # The model validates the enum, so the constraint is the second line of + # defence — it has to be exercised underneath the repository. + with pytest.raises(asyncpg.IntegrityConstraintViolationError): + await repo.pool.execute( + "INSERT INTO jobs (id, status, partition) VALUES ('bad', 'BOGUS', 'p')", + ) + + async def test_deleting_the_uploader_keeps_the_job_history(self, repo, postgres_store): + user_id = await _user(postgres_store) + await repo.create_job(_job(user_id=user_id)) + + await postgres_store.pool.execute("DELETE FROM users WHERE id = $1", user_id) + + job = await repo.get_job("task-1") + assert job is not None + assert job.user_id is None + + +class TestQueries: + async def test_list_jobs_filters_scopes_and_orders_newest_first(self, repo, postgres_store): + user_id = await _user(postgres_store) + other_id = await _user(postgres_store, "other") + now = datetime.now(UTC) + await repo.create_job(_job(id="old", user_id=user_id, created_at=now - timedelta(hours=1))) + await repo.create_job(_job(id="new", user_id=user_id, created_at=now)) + await repo.create_job(_job(id="theirs", user_id=other_id)) + + mine = await repo.list_jobs(user_id=user_id) + + assert [j.id for j in mine] == ["new", "old"] + + async def test_list_jobs_active_excludes_terminal_jobs(self, repo): + await repo.create_job(_job(id="running", status=DocumentStatus.CHUNKING)) + await repo.create_job(_job(id="done", status=DocumentStatus.COMPLETED)) + + assert [j.id for j in await repo.list_jobs(status="active")] == ["running"] + + async def test_list_jobs_paginates(self, repo): + for i in range(5): + await repo.create_job(_job(id=f"t{i}", created_at=datetime.now(UTC) + timedelta(seconds=i))) + + page = await repo.list_jobs(offset=2, limit=2) + + assert [j.id for j in page] == ["t2", "t1"] + + async def test_count_by_status_rolls_up_the_table(self, repo): + await repo.create_job(_job(id="a", status=DocumentStatus.QUEUED)) + await repo.create_job(_job(id="b", status=DocumentStatus.COMPLETED)) + await repo.create_job(_job(id="c", status=DocumentStatus.COMPLETED)) + + assert await repo.count_by_status() == {"QUEUED": 1, "COMPLETED": 2} + + +class TestRetention: + async def test_purge_removes_aged_terminal_jobs_only(self, repo): + stale = datetime.now(UTC) - timedelta(days=30) + await repo.create_job(_job(id="stale", status=DocumentStatus.COMPLETED, completed_at=stale)) + await repo.create_job(_job(id="fresh", status=DocumentStatus.COMPLETED, completed_at=datetime.now(UTC))) + # An in-flight job is old but not settled — it must survive. + await repo.create_job(_job(id="running", status=DocumentStatus.SERIALIZING, created_at=stale)) + + purged = await repo.purge_terminal_jobs(older_than_seconds=7 * 24 * 3600, keep_last=1000) + + assert purged == 1 + assert {j.id for j in await repo.list_jobs()} == {"fresh", "running"} + + async def test_purge_caps_the_table_when_jobs_settle_faster_than_the_ttl(self, repo): + now = datetime.now(UTC) + for i in range(5): + await repo.create_job( + _job(id=f"t{i}", status=DocumentStatus.COMPLETED, completed_at=now + timedelta(seconds=i)) + ) + + purged = await repo.purge_terminal_jobs(older_than_seconds=7 * 24 * 3600, keep_last=2) + + assert purged == 3 + assert {j.id for j in await repo.list_jobs()} == {"t4", "t3"} + + async def test_purge_of_an_empty_table_is_a_noop(self, repo): + assert await repo.purge_terminal_jobs(older_than_seconds=60, keep_last=10) == 0 + + +async def postgres_count(repo: PgJobRepository) -> int: + return await repo.pool.fetchval("SELECT COUNT(*)::int FROM jobs") diff --git a/tests/integration/repos/test_partition_repo.py b/tests/integration/repos/test_partition_repo.py index 8594e502e..a8b7d6533 100644 --- a/tests/integration/repos/test_partition_repo.py +++ b/tests/integration/repos/test_partition_repo.py @@ -69,14 +69,29 @@ async def test_delete_cascades_files_and_decrements_uploader_count( """Regression: ``files.partition_name`` has no DB-level CASCADE, so the repo must delete file rows itself before dropping the partition. Also verifies the per-uploader ``file_count`` decrement. + + The two slots are *reserved* first, because since #664 ``file_count`` is + a reserved+completed counter: admission charges the slot and + ``add_file_to_partition`` only consumes it. Without the reservations the + count would start at 0 and this would assert a vacuous 0 → 0. """ partition_repo = postgres_store.partition_repo document_repo = postgres_store.document_repo user_repo = postgres_store.user_repo - uploader = await user_repo.create_legacy_user(display_name="Uploader") + uploader = await user_repo.create_legacy_user( + display_name="Uploader", + external_user_id=None, + email=None, + is_admin=False, + file_quota=None, + ) uploader_id = uploader["id"] + # Admission: two uploads charged against an unlimited quota. + assert await user_repo.try_reserve_file_slot(uploader_id, default_quota=-1) == 1 + assert await user_repo.try_reserve_file_slot(uploader_id, default_quota=-1) == 2 + await partition_repo.create_partition("cascade-me") await document_repo.add_file_to_partition( file_id="f1", @@ -89,6 +104,9 @@ async def test_delete_cascades_files_and_decrements_uploader_count( user_id=uploader_id, ) assert await partition_repo.get_partition_file_count("cascade-me") == 2 + # Consumed, not double-counted: the inserts must not have re-incremented. + before = await user_repo.get_user_dict_by_id(uploader_id) + assert before["file_count"] == 2 assert await partition_repo.delete_partition("cascade-me") is True diff --git a/tests/integration/repos/test_user_repo_quota_reserve.py b/tests/integration/repos/test_user_repo_quota_reserve.py new file mode 100644 index 000000000..8e1d49f44 --- /dev/null +++ b/tests/integration/repos/test_user_repo_quota_reserve.py @@ -0,0 +1,134 @@ +"""Issue #664 — atomic quota reserve/release against a real Postgres. + +The point of the reserve is that admission is a *single* conditional +UPDATE, so N concurrent admits at quota can never overshoot. That +property is only meaningful against a real database, hence this suite +rather than a unit test with a fake pool. +""" + +from __future__ import annotations + +import asyncio + +import pytest +from core.models.user import User +from services.storage.postgres_store import PostgresStore + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +def _user(**overrides) -> User: + defaults = {"display_name": "Quota User", "is_admin": False} + defaults.update(overrides) + return User(**defaults) + + +async def _file_count(store: PostgresStore, user_id: int) -> int: + user = await store.user_repo.get_user(user_id) + assert user is not None + return user.file_count + + +class TestReserveSemantics: + async def test_reserve_under_quota_increments(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=3)) + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) == 1 + assert await _file_count(postgres_store, user.id) == 1 + + async def test_reserve_at_quota_rejects_and_does_not_increment(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=1)) + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) == 1 + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) is None + assert await _file_count(postgres_store, user.id) == 1 + + async def test_zero_quota_rejects_immediately(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=0)) + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=10) is None + assert await _file_count(postgres_store, user.id) == 0 + + async def test_null_quota_falls_back_to_default(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=None)) + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=1) == 1 + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=1) is None + + async def test_null_quota_with_negative_default_is_unlimited(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=None)) + for expected in (1, 2, 3): + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) == expected + + async def test_negative_per_user_quota_is_unlimited(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=-1)) + for expected in (1, 2, 3): + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=0) == expected + + async def test_explicit_per_user_quota_honored_when_default_is_negative(self, postgres_store: PostgresStore): + """A negative *global default* must not make a capped user unlimited.""" + user = await postgres_store.user_repo.create_user(_user(file_quota=2)) + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) == 1 + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) == 2 + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) is None + + async def test_admin_bypasses_quota(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(is_admin=True, file_quota=0)) + # Admins still get counted (file_count stays a truthful total) but are + # never rejected. + for expected in (1, 2, 3): + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=0) == expected + + async def test_unknown_user_returns_none(self, postgres_store: PostgresStore): + assert await postgres_store.user_repo.try_reserve_file_slot(99999, default_quota=-1) is None + + +class TestRelease: + async def test_release_decrements(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=5)) + await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) + await postgres_store.user_repo.release_file_slot(user.id) + assert await _file_count(postgres_store, user.id) == 0 + + async def test_release_clamps_at_zero(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=5)) + await postgres_store.user_repo.release_file_slot(user.id) + await postgres_store.user_repo.release_file_slot(user.id) + assert await _file_count(postgres_store, user.id) == 0 + + async def test_release_reopens_the_gate(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=1)) + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) == 1 + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) is None + await postgres_store.user_repo.release_file_slot(user.id) + assert await postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) == 1 + + async def test_release_on_unknown_user_is_a_noop(self, postgres_store: PostgresStore): + await postgres_store.user_repo.release_file_slot(99999) + + +class TestConcurrency: + @pytest.mark.parametrize("quota", [1, 5]) + async def test_parallel_admits_never_overshoot(self, postgres_store: PostgresStore, quota: int): + """The regression this issue is about: burst admission at quota. + + 20 concurrent reserves against a quota of ``quota`` must grant + exactly ``quota`` slots — never more — and leave ``file_count`` + exactly at the quota. + """ + concurrency = 20 + user = await postgres_store.user_repo.create_user(_user(file_quota=quota)) + + results = await asyncio.gather( + *(postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) for _ in range(concurrency)) + ) + + granted = [r for r in results if r is not None] + assert len(granted) == quota, f"expected exactly {quota} admits, got {len(granted)}: {results}" + # Every granted reservation reports a distinct, contiguous count. + assert sorted(granted) == list(range(1, quota + 1)) + assert await _file_count(postgres_store, user.id) == quota + + async def test_parallel_reserve_and_release_settles(self, postgres_store: PostgresStore): + user = await postgres_store.user_repo.create_user(_user(file_quota=10)) + await asyncio.gather( + *(postgres_store.user_repo.try_reserve_file_slot(user.id, default_quota=-1) for _ in range(10)) + ) + await asyncio.gather(*(postgres_store.user_repo.release_file_slot(user.id) for _ in range(10))) + assert await _file_count(postgres_store, user.id) == 0 diff --git a/tests/unit/api/dependencies/test_auth.py b/tests/unit/api/dependencies/test_auth.py index 58fba6719..0995820f7 100644 --- a/tests/unit/api/dependencies/test_auth.py +++ b/tests/unit/api/dependencies/test_auth.py @@ -1,13 +1,15 @@ +import asyncio from types import SimpleNamespace import pytest from api.dependencies.auth import ( check_user_file_quota, + commit_quota_reservation, ensure_partition_role, require_partitions_viewer, require_task_owner, ) -from core.utils.exceptions import AuthError +from core.utils.exceptions import AuthError, OpenRAGError from fastapi import HTTPException from services.orchestrators.auth_service import AuthService @@ -231,57 +233,162 @@ async def test_require_task_owner_rejects_non_owner_non_admin(): assert exc.value.status_code == 403 +class RecordingAuthService(FakeAuthService): + """Stands in for AuthService's reserve/release pair (#664).""" + + def __init__(self, *, grant: bool = True, raise_on_release: bool = False) -> None: + self.grant = grant + self.raise_on_release = raise_on_release + self.reserved: list[tuple[int, int]] = [] + self.released: list[int] = [] + + async def reserve_file_slot(self, user_id: int, *, default_quota: int) -> int: + self.reserved.append((user_id, default_quota)) + if not self.grant: + raise OpenRAGError( + "File quota exceeded.", + code="FILE_QUOTA_EXCEEDED", + status_code=403, + ) + return len(self.reserved) + + async def release_file_slot(self, user_id: int) -> None: + self.released.append(user_id) + + +async def _drive_quota_dep(auth_service, user, default_quota: int, *, commit: bool, boom: Exception | None = None): + """Run the dependency generator the way FastAPI's exit stack does.""" + gen = check_user_file_quota( + user=user, + auth_service=auth_service, + config=_config(default_file_quota=default_quota), + ) + reservation = await gen.__anext__() + if commit: + commit_quota_reservation(reservation) + try: + if boom is not None: + # FastAPI throws the endpoint's exception back into the generator. + await gen.athrow(boom) + else: + await gen.__anext__() + except StopAsyncIteration: + pass + except type(boom) if boom is not None else (): + pass + return reservation + + @pytest.mark.asyncio -async def test_check_user_file_quota_reads_pending_count_through_job_service(): - job_service = FakeJobService(pending_count=2) +async def test_check_user_file_quota_reserves_a_slot_atomically(): + auth_service = RecordingAuthService() + + reservation = await _drive_quota_dep( + auth_service, + {"id": 7, "file_count": 1, "file_quota": 5}, + default_quota=10, + commit=True, + ) - user = await check_user_file_quota( - user={"id": 7, "file_count": 1, "file_quota": 5}, - auth_service=FakeAuthService, - job_service=job_service, + assert auth_service.reserved == [(7, 10)] + assert reservation.user_id == 7 + assert reservation.committed is True + + +@pytest.mark.asyncio +async def test_check_user_file_quota_rejects_with_403_when_no_slot(): + """No row from the conditional UPDATE ⇒ the upload is refused.""" + auth_service = RecordingAuthService(grant=False) + + gen = check_user_file_quota( + user={"id": 8, "file_count": 5, "file_quota": 5}, + auth_service=auth_service, config=_config(default_file_quota=10), ) + with pytest.raises(HTTPException) as exc: + await gen.__anext__() - assert user["id"] == 7 - assert job_service.pending_checks == [7] + assert exc.value.status_code == 403 + assert "quota" in exc.value.detail.lower() + # A rejected reserve took nothing, so there is nothing to hand back. + assert auth_service.released == [] @pytest.mark.asyncio -async def test_check_user_file_quota_skips_pending_count_when_default_is_unlimited(): - """Skip queue I/O when the resolved quota is unlimited.""" - job_service = FakeJobService(pending_count=2) +async def test_committed_reservation_is_not_released(): + """After dispatch the worker owns the slot; teardown must keep its hands off.""" + auth_service = RecordingAuthService() - user = await check_user_file_quota( - user={"id": 7, "file_count": 1, "file_quota": None}, - auth_service=FakeAuthService, - job_service=job_service, - config=_config(default_file_quota=-1), + await _drive_quota_dep(auth_service, {"id": 7}, default_quota=5, commit=True) + + assert auth_service.released == [] + + +@pytest.mark.asyncio +async def test_uncommitted_reservation_is_released_on_teardown(): + """The route returned early (e.g. 409 duplicate) — the slot goes back.""" + auth_service = RecordingAuthService() + + await _drive_quota_dep(auth_service, {"id": 7}, default_quota=5, commit=False) + + assert auth_service.released == [7] + + +@pytest.mark.asyncio +async def test_reservation_is_released_when_the_endpoint_raises(): + """FastAPI throws the endpoint error into the dependency; cleanup still runs.""" + auth_service = RecordingAuthService() + + await _drive_quota_dep( + auth_service, + {"id": 7}, + default_quota=5, + commit=False, + boom=HTTPException(status_code=409, detail="already exists"), ) - assert user["id"] == 7 - assert job_service.pending_checks == [] + assert auth_service.released == [7] @pytest.mark.asyncio -async def test_check_user_file_quota_default_zero_enforces_zero(): - allowed_job_service = FakeJobService(pending_count=0) - user = await check_user_file_quota( - user={"id": 7, "file_count": 0, "file_quota": 1}, - auth_service=EnforcingAuthService, - job_service=allowed_job_service, - config=_config(default_file_quota=0), +async def test_reservation_is_released_when_the_client_disconnects(): + """A cancelled request must not leak the slot it admitted.""" + auth_service = RecordingAuthService() + + await _drive_quota_dep( + auth_service, + {"id": 7}, + default_quota=5, + commit=False, + boom=asyncio.CancelledError(), ) - assert user["id"] == 7 - assert allowed_job_service.pending_checks == [7] - denied_job_service = FakeJobService(pending_count=0) - with pytest.raises(HTTPException) as exc: - await check_user_file_quota( - user={"id": 8, "file_count": 1, "file_quota": None}, - auth_service=EnforcingAuthService, - job_service=denied_job_service, - config=_config(default_file_quota=0), - ) - assert exc.value.status_code == 403 - assert exc.value.detail == "File quota exceeded" - assert denied_job_service.pending_checks == [8] + assert auth_service.released == [7] + + +@pytest.mark.asyncio +async def test_check_user_file_quota_no_op_without_a_user_id(): + """Nothing durable to charge (auth disabled) ⇒ no reserve, no release.""" + auth_service = RecordingAuthService() + + reservation = await _drive_quota_dep(auth_service, {}, default_quota=5, commit=False) + + assert reservation is None + assert auth_service.reserved == [] + assert auth_service.released == [] + + +@pytest.mark.asyncio +async def test_admin_reservation_still_counts_but_is_never_rejected(): + """Admins bypass the *limit*, not the counter — the SQL predicate decides.""" + auth_service = RecordingAuthService() + + await _drive_quota_dep(auth_service, {"id": 1, "is_admin": True}, default_quota=0, commit=True) + + assert auth_service.reserved == [(1, 0)] + + +def test_commit_quota_reservation_ignores_non_reservations(): + """Tests routinely override the dependency with a plain stub.""" + commit_quota_reservation(None) + commit_quota_reservation({"id": 1}) diff --git a/tests/unit/api/routers/admin/test_indexing_quota_release.py b/tests/unit/api/routers/admin/test_indexing_quota_release.py new file mode 100644 index 000000000..35b1c82bb --- /dev/null +++ b/tests/unit/api/routers/admin/test_indexing_quota_release.py @@ -0,0 +1,314 @@ +"""Issue #664 — every pre-dispatch exit must hand the quota slot back. + +``check_user_file_quota`` reserves a slot *before* the route body runs, so +from that point on any early return is a path where an admitted upload never +becomes a file. Each one is covered here end-to-end through the real +dependency (not a stub), because a leak here is silent: it costs the user a +slot forever and the only symptom is a quota that mysteriously shrinks. +""" + +from __future__ import annotations + +import io +from types import SimpleNamespace + +import httpx +import pytest +from api.dependencies.auth import current_user, require_partition_editor +from api.dependencies.files import validate_file_format, validate_file_id, validate_metadata +from api.error_handlers import register_error_handlers +from api.routers.admin.indexing import router as indexer_router +from core.utils.exceptions import OpenRAGError, PartitionNotFoundError +from di.providers import get_auth_service, get_config, get_indexing_service +from fastapi import FastAPI, UploadFile + +USER = {"id": 42, "is_admin": False, "file_quota": 5, "file_count": 0} + + +class RecordingAuthService: + """Tracks the reserve/release pair the real AuthService would perform.""" + + def __init__(self, *, grant: bool = True) -> None: + self.grant = grant + self.reserved: list[int] = [] + self.released: list[int] = [] + + async def reserve_file_slot(self, user_id: int, *, default_quota: int) -> int: + if not self.grant: + raise OpenRAGError("File quota exceeded.", code="FILE_QUOTA_EXCEEDED", status_code=403) + self.reserved.append(user_id) + return len(self.reserved) + + async def release_file_slot(self, user_id: int) -> None: + self.released.append(user_id) + + # copy_file's source-partition check goes through these. + @staticmethod + def check_partition_access(**kwargs) -> bool: + return True + + +class FakeIndexingService: + def __init__( + self, + *, + exists: bool = False, + workspace: dict | None = None, + add_error: Exception | None = None, + copy_result: bool = True, + copy_error: Exception | None = None, + ) -> None: + self.exists = exists + self.workspace = workspace + self.add_error = add_error + self.copy_result = copy_result + self.copy_error = copy_error + self.dispatched = 0 + self.add_kwargs: dict = {} + + async def file_exists(self, file_id: str, partition: str) -> bool: + return self.exists + + async def get_workspace(self, workspace_id: str): + return self.workspace + + async def add_file(self, **kwargs): + if self.add_error is not None: + raise self.add_error + self.add_kwargs = kwargs + self.dispatched += 1 + return "task-1" + + async def copy_file(self, **kwargs) -> bool: + if self.copy_error is not None: + raise self.copy_error + return self.copy_result + + +def _no_metadata() -> dict: + return {} + + +def _build_app(tmp_path, auth_service, service, *, content: bytes = b"hi"): + app = FastAPI() + register_error_handlers(app) + app.include_router(indexer_router, prefix="/indexer") + + cfg = SimpleNamespace( + paths=SimpleNamespace(data_dir=str(tmp_path / "data")), + rdb=SimpleNamespace(default_file_quota=5), + server=SimpleNamespace(preferred_url_scheme="http"), + ) + + app.dependency_overrides[validate_file_id] = lambda: "f1" + app.dependency_overrides[validate_file_format] = lambda: UploadFile(file=io.BytesIO(content), filename="doc.txt") + app.dependency_overrides[validate_metadata] = _no_metadata + app.dependency_overrides[require_partition_editor] = lambda: USER + app.dependency_overrides[current_user] = lambda: USER + app.dependency_overrides[get_auth_service] = lambda: auth_service + app.dependency_overrides[get_config] = lambda: cfg + app.dependency_overrides[get_indexing_service] = lambda: service + # check_user_file_quota itself is intentionally NOT overridden. + return app + + +async def _post(app, url, **kwargs): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + return await client.post(url, **kwargs) + + +# ── add_file ─────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_successful_dispatch_keeps_the_slot(tmp_path): + auth = RecordingAuthService() + service = FakeIndexingService() + + resp = await _post(_build_app(tmp_path, auth, service), "/indexer/partition/p1/file/f1", data={"_": "1"}) + + assert resp.status_code == 201 + assert service.dispatched == 1 + assert auth.reserved == [42] + assert auth.released == [] + + +@pytest.mark.asyncio +async def test_dispatch_tells_the_worker_the_slot_is_already_reserved(tmp_path): + """The handoff is a *signal*, and nothing else pins that it is sent. + + ``commit_quota_reservation`` only stops the request's teardown from + releasing; it is ``quota_reserved=True`` that makes the worker take + ownership and release on failure/cancellation. Drop the flag and the two + halves disagree: teardown declines to release because the reservation was + committed, and the worker declines because it was never told it owns one — + so every failed upload leaks a slot permanently, in silence. + + The rest of this module asserts the release paths; this asserts the wire + between them, which the release tests cannot see because they never reach + a worker. + """ + service = FakeIndexingService() + + resp = await _post( + _build_app(tmp_path, RecordingAuthService(), service), + "/indexer/partition/p1/file/f1", + data={"_": "1"}, + ) + + assert resp.status_code == 201 + assert service.add_kwargs.get("quota_reserved") is True + + +@pytest.mark.asyncio +async def test_quota_exceeded_rejects_before_touching_the_route(tmp_path): + auth = RecordingAuthService(grant=False) + service = FakeIndexingService() + + resp = await _post(_build_app(tmp_path, auth, service), "/indexer/partition/p1/file/f1", data={"_": "1"}) + + assert resp.status_code == 403 + assert service.dispatched == 0 + assert auth.released == [] + + +@pytest.mark.asyncio +async def test_duplicate_file_409_releases_the_slot(tmp_path): + auth = RecordingAuthService() + service = FakeIndexingService(exists=True) + + resp = await _post(_build_app(tmp_path, auth, service), "/indexer/partition/p1/file/f1", data={"_": "1"}) + + assert resp.status_code == 409 + assert auth.released == [42] + + +@pytest.mark.asyncio +async def test_oversize_upload_releases_the_slot(tmp_path, monkeypatch): + monkeypatch.setattr("api.dependencies.files._max_upload_size_bytes", lambda: 8) + auth = RecordingAuthService() + service = FakeIndexingService() + app = _build_app(tmp_path, auth, service, content=b"x" * 100) + + resp = await _post(app, "/indexer/partition/p1/file/f1", data={"_": "1"}) + + assert resp.status_code == 413 + assert auth.released == [42] + + +@pytest.mark.asyncio +async def test_bad_workspace_ids_400_releases_the_slot(tmp_path): + auth = RecordingAuthService() + service = FakeIndexingService() + + resp = await _post( + _build_app(tmp_path, auth, service), + "/indexer/partition/p1/file/f1", + data={"workspace_ids": "not-json"}, + ) + + assert resp.status_code == 400 + assert auth.released == [42] + + +@pytest.mark.asyncio +async def test_unknown_workspace_404_releases_the_slot(tmp_path): + auth = RecordingAuthService() + service = FakeIndexingService(workspace=None) + + resp = await _post( + _build_app(tmp_path, auth, service), + "/indexer/partition/p1/file/f1", + data={"workspace_ids": '["ws-1"]'}, + ) + + assert resp.status_code == 404 + assert auth.released == [42] + + +@pytest.mark.asyncio +async def test_dispatch_failure_releases_the_slot(tmp_path): + """Anything raising between admission and a queued job gives the slot back.""" + auth = RecordingAuthService() + service = FakeIndexingService(add_error=PartitionNotFoundError("Partition 'p1' does not exist.")) + + resp = await _post(_build_app(tmp_path, auth, service), "/indexer/partition/p1/file/f1", data={"_": "1"}) + + assert resp.status_code == 404 + assert service.dispatched == 0 + assert auth.released == [42] + + +# ── copy_file ────────────────────────────────────────────────────────────── + + +def _copy_app(tmp_path, auth, service): + from api.dependencies.auth import current_user_partitions + from di.providers import get_partition_service + + app = _build_app(tmp_path, auth, service) + app.dependency_overrides[current_user_partitions] = lambda: [{"partition": "src", "role": "owner"}] + app.dependency_overrides[get_partition_service] = lambda: SimpleNamespace() + return app + + +@pytest.mark.asyncio +async def test_copy_file_is_quota_gated_and_keeps_the_slot_on_success(tmp_path): + auth = RecordingAuthService() + service = FakeIndexingService(copy_result=True) + + resp = await _post( + _copy_app(tmp_path, auth, service), + "/indexer/partition/p1/file/f1/copy", + data={"source_partition": "src", "source_file_id": "src-f1"}, + ) + + assert resp.status_code == 201 + assert auth.reserved == [42] + assert auth.released == [] + + +@pytest.mark.asyncio +async def test_copy_file_over_quota_is_rejected(tmp_path): + """Regression: copy_file is the second quota-gated route (#664).""" + auth = RecordingAuthService(grant=False) + + resp = await _post( + _copy_app(tmp_path, auth, FakeIndexingService()), + "/indexer/partition/p1/file/f1/copy", + data={"source_partition": "src", "source_file_id": "src-f1"}, + ) + + assert resp.status_code == 403 + + +@pytest.mark.asyncio +async def test_copy_file_that_creates_no_row_releases_the_slot(tmp_path): + """Empty source, or the target already existed — no file, no slot.""" + auth = RecordingAuthService() + service = FakeIndexingService(copy_result=False) + + resp = await _post( + _copy_app(tmp_path, auth, service), + "/indexer/partition/p1/file/f1/copy", + data={"source_partition": "src", "source_file_id": "src-f1"}, + ) + + assert resp.status_code == 201 + assert auth.released == [42] + + +@pytest.mark.asyncio +async def test_copy_file_error_releases_the_slot(tmp_path): + auth = RecordingAuthService() + service = FakeIndexingService(copy_error=RuntimeError("milvus down")) + + with pytest.raises(RuntimeError, match="milvus down"): + await _post( + _copy_app(tmp_path, auth, service), + "/indexer/partition/p1/file/f1/copy", + data={"source_partition": "src", "source_file_id": "src-f1"}, + ) + + assert auth.released == [42] diff --git a/tests/unit/core/utils/test_text.py b/tests/unit/core/utils/test_text.py index 005867f32..b410f1acb 100644 --- a/tests/unit/core/utils/test_text.py +++ b/tests/unit/core/utils/test_text.py @@ -74,3 +74,25 @@ def get_num_tokens(self, _text: str) -> int: assert length_function("hello") == 1 assert "enable_thinking" not in captured + + +def test_truncate_error_text_keeps_the_tail_not_the_head(): + """Which end survives is the whole point, and nothing pinned it. + + A Python traceback puts the exception type and message *last*, so a + head-truncating implementation would retain the same dispatcher frames + on every task and discard the only line an operator needs. Assertions on + the length and the marker alone are satisfied by that implementation too, + which is why this asserts on the surviving end specifically. + """ + from core.utils.text import truncate_error_text + + head = "OUTERMOST_FRAME_MARKER" + tail = "ValueError: the actual cause" + tb = head + (' File "x.py", line 1, in f\n' * 2000) + tail + + out = truncate_error_text(tb, 200) + + assert out.endswith(tail), "the exception message must survive truncation" + assert out.endswith(tb[-200:]), "the retained slice must be the trailing one" + assert head not in out, "the outermost frames are the part to drop" diff --git a/tests/unit/services/orchestrators/test_indexing_service.py b/tests/unit/services/orchestrators/test_indexing_service.py index cb5ccf093..3615ae29b 100644 --- a/tests/unit/services/orchestrators/test_indexing_service.py +++ b/tests/unit/services/orchestrators/test_indexing_service.py @@ -53,9 +53,11 @@ async def dispatch_indexing( embedder_name=None, require_existing_partition=False, allow_legacy_require_existing_partition_retry=False, + quota_reserved=False, ): self.dispatched.append( { + "quota_reserved": quota_reserved, "path": path, "metadata": metadata, "partition": partition, diff --git a/tests/unit/services/orchestrators/test_job_service.py b/tests/unit/services/orchestrators/test_job_service.py index 5f8533b6d..43518ef74 100644 --- a/tests/unit/services/orchestrators/test_job_service.py +++ b/tests/unit/services/orchestrators/test_job_service.py @@ -131,3 +131,251 @@ async def test_get_user_pending_task_count_uses_task_state_manager(): pending = await JobService(FakeTSM(info=info)).get_user_pending_task_count(7) assert pending == 2 + + +# --------------------------------------------------------------------------- +# Durable reads (issue #660) — Postgres is the source of truth, the actor is +# a hot cache whose terminal entries are evicted. +# --------------------------------------------------------------------------- + + +class FakeJobRepo: + def __init__(self, jobs=None, counts=None, pending=0, boom=False): + self._jobs = jobs or [] + self._counts = counts or {} + self._pending = pending + self._boom = boom + self.list_calls: list[dict] = [] + + def _check(self): + if self._boom: + raise RuntimeError("postgres down") + + async def list_jobs(self, status=None, offset=0, limit=50, user_id=None): + self._check() + self.list_calls.append({"status": status, "offset": offset, "limit": limit, "user_id": user_id}) + return list(self._jobs)[offset : offset + limit] + + async def get_job(self, job_id): + self._check() + return next((j for j in self._jobs if j.id == job_id), None) + + async def count_by_status(self): + self._check() + return dict(self._counts) + + +def _job(**kwargs): + from core.models.catalog import IndexationJob + + base = {"id": "t1", "partition": "p", "file_id": "f1", "user_id": 7, "job_metadata": {"filename": "a.pdf"}} + base.update(kwargs) + return IndexationJob(**base) + + +@pytest.mark.asyncio +async def test_get_queue_info_counts_come_from_postgres_when_available(): + repo = FakeJobRepo(counts={"QUEUED": 2, "COMPLETED": 5, "FAILED": 1}) + svc = JobService(task_state_manager=FakeTSM(states={"zombie": "QUEUED"}), job_repo=repo) + + out = await svc.get_queue_info() + + assert out["tasks"]["active"] == 2 + assert out["tasks"]["total_completed"] == 5 + assert out["tasks"]["total_failed"] == 1 + # pool info still comes from the actor + assert out["workers"]["pool_size"] == 2 + + +@pytest.mark.asyncio +async def test_get_queue_info_falls_back_to_the_actor_when_postgres_is_down(): + svc = JobService(task_state_manager=FakeTSM(states={"a": "QUEUED"}), job_repo=FakeJobRepo(boom=True)) + + out = await svc.get_queue_info() + + assert out["tasks"]["active"] == 1 + + +@pytest.mark.asyncio +async def test_list_tasks_reads_durable_jobs_and_scopes_non_admins(): + from core.models.catalog import DocumentStatus + + repo = FakeJobRepo(jobs=[_job(status=DocumentStatus.COMPLETED)]) + svc = JobService(task_state_manager=FakeTSM(), job_repo=repo) + + rows = await svc.list_tasks(is_admin=False, user_id=7, task_status="completed") + + assert rows == [ + { + "task_id": "t1", + "state": "COMPLETED", + "details": {"file_id": "f1", "partition": "p", "metadata": {"filename": "a.pdf"}, "user_id": 7}, + } + ] + assert repo.list_calls[0]["user_id"] == 7 + assert repo.list_calls[0]["status"] == "completed" + + +class _WarningRecorder: + """Stands in for the module logger: loguru bypasses caplog and capsys.""" + + def __init__(self) -> None: + self.warnings: list[str] = [] + + def warning(self, message, **kwargs): + self.warnings.append(message) + + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + +def _recording_logger(monkeypatch) -> _WarningRecorder: + from services.orchestrators import job_service as job_service_module + + recorder = _WarningRecorder() + monkeypatch.setattr(job_service_module, "logger", recorder) + return recorder + + +@pytest.mark.asyncio +async def test_list_tasks_caps_the_page_and_warns_when_it_truncates(monkeypatch): + """A capped page must not be handed back as if it were the whole queue.""" + from services.orchestrators.job_service import _LIST_LIMIT + + recorder = _recording_logger(monkeypatch) + repo = FakeJobRepo(jobs=[_job(id=f"t{i}") for i in range(_LIST_LIMIT + 50)]) + svc = JobService(task_state_manager=FakeTSM(), job_repo=repo) + + rows = await svc.list_tasks(is_admin=True, user_id=1) + + assert len(rows) == _LIST_LIMIT + # asked for one more than it returns: that extra row is how truncation is seen + assert repo.list_calls[0]["limit"] == _LIST_LIMIT + 1 + assert any("truncated" in w for w in recorder.warnings) + + +@pytest.mark.asyncio +async def test_list_tasks_does_not_warn_on_an_exactly_full_page(monkeypatch): + """Exactly _LIST_LIMIT jobs is a complete answer, not a truncated one.""" + from services.orchestrators.job_service import _LIST_LIMIT + + recorder = _recording_logger(monkeypatch) + repo = FakeJobRepo(jobs=[_job(id=f"t{i}") for i in range(_LIST_LIMIT)]) + svc = JobService(task_state_manager=FakeTSM(), job_repo=repo) + + rows = await svc.list_tasks(is_admin=True, user_id=1) + + assert len(rows) == _LIST_LIMIT + assert recorder.warnings == [] + + +@pytest.mark.asyncio +async def test_list_tasks_does_not_scope_admins_to_their_own_jobs(): + repo = FakeJobRepo(jobs=[]) + await JobService(task_state_manager=FakeTSM(), job_repo=repo).list_tasks(is_admin=True, user_id=1) + + assert repo.list_calls[0]["user_id"] is None + + +@pytest.mark.asyncio +async def test_list_tasks_falls_back_to_the_actor_when_postgres_is_down(): + info = {"t1": {"state": "FAILED", "details": {"file_id": "f"}, "user": 1}} + svc = JobService(task_state_manager=FakeTSM(info=info), job_repo=FakeJobRepo(boom=True)) + + rows = await svc.list_tasks(is_admin=True, user_id=1) + + assert rows == [{"task_id": "t1", "state": "FAILED", "details": {"file_id": "f"}}] + + +@pytest.mark.asyncio +async def test_get_task_details_survives_a_restart_via_postgres(): + svc = JobService(task_state_manager=FakeTSM(), job_repo=FakeJobRepo(jobs=[_job()])) + + assert await svc.get_task_details("t1") == { + "file_id": "f1", + "partition": "p", + "metadata": {"filename": "a.pdf"}, + "user_id": 7, + } + + +@pytest.mark.asyncio +async def test_get_user_pending_task_count_stays_on_the_in_memory_cache(): + """The quota gate must not be jammable by an orphaned job row. + + A durable job leaves the active states only when a worker writes a terminal + transition, so a crash mid-dispatch would hold the user's quota open-endedly + (retention sweeps terminal rows only). See the method docstring / #664. + """ + info = {"t1": {"state": "QUEUED", "details": {}, "user_id": 7}} + repo = FakeJobRepo(pending=99) + svc = JobService(task_state_manager=FakeTSM(info=info), job_repo=repo) + + assert await svc.get_user_pending_task_count(7) == 1 + + +@pytest.mark.asyncio +async def test_list_tasks_falls_back_to_the_cache_when_postgres_has_no_rows(): + """An empty ``jobs`` table is a durable *miss*, not an empty queue. + + The actor is detached, so it outlives the API restart that first deploys the + durable store: every task dispatched before the cutover has no row. Answering + ``[]`` authoritatively hides work that is actively indexing — and the same + happens for any task whose best-effort ``create_job`` was swallowed. + """ + info = {"t1": {"state": "QUEUED", "details": {"file_id": "f"}, "user": 1}} + repo = FakeJobRepo(jobs=[]) + svc = JobService(task_state_manager=FakeTSM(info=info), job_repo=repo) + + rows = await svc.list_tasks(is_admin=True, user_id=1) + + assert rows == [{"task_id": "t1", "state": "QUEUED", "details": {"file_id": "f"}}] + assert repo.list_calls, "the durable store should still be consulted first" + + +@pytest.mark.asyncio +async def test_get_queue_info_falls_back_to_the_cache_when_postgres_has_no_rows(): + """Same miss-vs-empty distinction for the roll-up: never report 0 active + while the cache is holding live tasks.""" + svc = JobService( + task_state_manager=FakeTSM(states={"a": "QUEUED", "b": "SERIALIZING"}), + job_repo=FakeJobRepo(counts={}), + ) + + tasks = (await svc.get_queue_info())["tasks"] + + assert tasks["active"] == 2 + assert tasks["active_statuses"] == {"QUEUED": 1, "SERIALIZING": 1, "CHUNKING": 0, "INSERTING": 0} + + +@pytest.mark.asyncio +async def test_a_populated_postgres_still_wins_over_the_cache(): + """The fallback is for misses only — it must not resurrect evicted entries. + + Guards the obvious over-correction: falling back whenever the durable result + looks small would let the actor's stale copy shadow the source of truth. + """ + repo = FakeJobRepo(jobs=[_job(id="durable")]) + svc = JobService( + task_state_manager=FakeTSM(info={"stale": {"state": "QUEUED", "details": {}, "user": 1}}), + job_repo=repo, + ) + + rows = await svc.list_tasks(is_admin=True, user_id=1) + + assert [r["task_id"] for r in rows] == ["durable"] + + +@pytest.mark.asyncio +async def test_list_tasks_fails_closed_for_an_anonymous_non_admin(): + """list_jobs(user_id=None) means *every* job — never hand that to a user. + + The durable read scopes with user_id=None if is_admin else user_id, so a + non-admin arriving without an id would select the whole table. Guard here + rather than trust every caller to have resolved one. + """ + repo = FakeJobRepo(jobs=[_job(id="t1", user_id=1), _job(id="t2", user_id=2)]) + svc = JobService(task_state_manager=FakeTSM(), job_repo=repo) + + assert await svc.list_tasks(is_admin=False, user_id=None) == [] + assert repo.list_calls == [] diff --git a/tests/unit/services/orchestrators/test_mcp_service.py b/tests/unit/services/orchestrators/test_mcp_service.py index 6ab6a060d..f854787e9 100644 --- a/tests/unit/services/orchestrators/test_mcp_service.py +++ b/tests/unit/services/orchestrators/test_mcp_service.py @@ -12,7 +12,7 @@ import httpx import pytest -from core.utils.exceptions import ValidationError +from core.utils.exceptions import OpenRAGError, ValidationError from services.orchestrators.mcp_service import MCPService # --------------------------------------------------------------------------- @@ -87,9 +87,10 @@ async def create_partition(self, partition, user_id, *, max_owned=None): class FakeIndexing: - def __init__(self, *, state="COMPLETED", error="boom"): + def __init__(self, *, state="COMPLETED", error="boom", copy_created=True): self._state = state self._error = error + self.copy_created = copy_created self.deleted: list[tuple[str, str]] = [] self.updated: list[tuple] = [] self.copied: list[dict] = [] @@ -109,12 +110,39 @@ async def update_metadata(self, file_id, metadata, partition, user): async def copy_file(self, **kwargs): self.copied.append(kwargs) + return self.copy_created async def add_file(self, **kwargs): self.added.append(kwargs) return "task-123" +class FakeAuth: + """Records reserve/release so tests can assert the slot is settled. + + Mirrors ``AuthService``: reserving over quota raises ``FILE_QUOTA_EXCEEDED`` + (403), releasing never raises. + """ + + def __init__(self, *, over_quota: bool = False) -> None: + self.over_quota = over_quota + self.reserved: list[int] = [] + self.released: list[int] = [] + + async def reserve_file_slot(self, user_id: int, *, default_quota: int) -> int: + if self.over_quota: + raise OpenRAGError( + "File quota exceeded.", + code="FILE_QUOTA_EXCEEDED", + status_code=403, + ) + self.reserved.append(user_id) + return len(self.reserved) + + async def release_file_slot(self, user_id: int) -> None: + self.released.append(user_id) + + class FakeJobs: def __init__(self, *, details=None, tasks=None): self._details = details @@ -153,7 +181,9 @@ def _service( indexing=None, jobs=None, conversion=None, + auth=None, vector_store=None, + default_file_quota=-1, default_top_k=5, max_top_k=50, similarity_threshold=0.8, @@ -164,8 +194,10 @@ def _service( indexing_service=indexing or FakeIndexing(), job_service=jobs or FakeJobs(), conversion_service=conversion or FakeConversion(), + auth_service=auth or FakeAuth(), vector_store=vector_store or FakeVectorStore(), collection="chunks", + default_file_quota=default_file_quota, default_top_k=default_top_k, max_top_k=max_top_k, similarity_threshold=similarity_threshold, @@ -834,3 +866,176 @@ async def test_get_file_chunks_caps_page_size(): out = await svc.get_file_chunks(partition="a", file_id="f1", allowed_partitions=["a"], offset=0, limit=-1) assert len(out["chunks"]) == _MAX_CHUNKS_PER_CALL assert out["has_more"] is True + + +# --------------------------------------------------------------------------- +# File-quota admission (issue #664) +# --------------------------------------------------------------------------- +# +# ``users.file_count`` is a reserved+completed counter and there is no +# completion-time increment any more, so a path that writes a ``files`` row +# without reserving is invisible to the quota — while the delete path still +# decrements unconditionally, driving the count *below* reality and handing out +# free slots. These two tools create rows, so they must reserve like the HTTP +# routes do in ``check_user_file_quota``. + + +class _CopySourceOnly(FakePartitions): + """Source file exists, destination does not — copy_file's happy path.""" + + async def file_exists(self, file_id, partition): + return file_id == "src" + + +async def _ok_download(url, dest): + dest.write_bytes(b"data") + + +def _url_svc(auth, indexing, *, partitions=None): + return _service( + auth=auth, + indexing=indexing, + partitions=partitions or FakePartitions(exists=False, partition_exists=False), + ) + + +@pytest.mark.asyncio +async def test_index_url_reserves_a_slot_and_hands_it_to_the_worker(monkeypatch): + auth, indexing = FakeAuth(), FakeIndexing() + svc = _url_svc(auth, indexing) + monkeypatch.setattr(svc, "_safe_download", _ok_download) + + out = await svc.index_url( + url="https://example.com/r.pdf", partition="p", file_id="f1", allowed_partitions=["all"], user_id=7 + ) + + assert out["task_id"] == "task-123" + assert auth.reserved == [7] + # Dispatched — the worker owns the slot now and must not have it taken back. + assert auth.released == [] + assert indexing.added[0]["quota_reserved"] is True + + +@pytest.mark.asyncio +async def test_index_url_over_quota_is_rejected_before_any_download(monkeypatch): + auth, indexing = FakeAuth(over_quota=True), FakeIndexing() + svc = _url_svc(auth, indexing) + + async def never(url, dest): + raise AssertionError("must not download when the quota already refused") + + monkeypatch.setattr(svc, "_safe_download", never) + + with pytest.raises(OpenRAGError) as exc: + await svc.index_url( + url="https://example.com/r.pdf", partition="p", file_id="f1", allowed_partitions=["all"], user_id=7 + ) + + assert exc.value.code == "FILE_QUOTA_EXCEEDED" + assert indexing.added == [] + assert auth.released == [] # the reserve refused; there is nothing to give back + + +@pytest.mark.asyncio +async def test_index_url_releases_the_slot_when_the_download_fails(monkeypatch): + auth, indexing = FakeAuth(), FakeIndexing() + svc = _url_svc(auth, indexing) + + async def boom(url, dest): + raise httpx.ConnectError("nope") + + monkeypatch.setattr(svc, "_safe_download", boom) + + with pytest.raises(RuntimeError): + await svc.index_url( + url="https://example.com/r.pdf", partition="p", file_id="f1", allowed_partitions=["all"], user_id=7 + ) + + assert auth.reserved == [7] and auth.released == [7] + assert indexing.added == [] + + +@pytest.mark.asyncio +async def test_index_url_releases_the_slot_when_dispatch_fails(monkeypatch): + class BoomIndexing(FakeIndexing): + async def add_file(self, **kwargs): + raise RuntimeError("pool unreachable") + + auth = FakeAuth() + svc = _url_svc(auth, BoomIndexing()) + monkeypatch.setattr(svc, "_safe_download", _ok_download) + + with pytest.raises(RuntimeError, match="pool unreachable"): + await svc.index_url( + url="https://example.com/r.pdf", partition="p", file_id="f1", allowed_partitions=["all"], user_id=7 + ) + + assert auth.released == [7] + + +@pytest.mark.asyncio +async def test_index_url_without_a_user_reserves_nothing(monkeypatch): + auth, indexing = FakeAuth(), FakeIndexing() + svc = _url_svc(auth, indexing) + monkeypatch.setattr(svc, "_safe_download", _ok_download) + + await svc.index_url( + url="https://example.com/r.pdf", partition="p", file_id="f1", allowed_partitions=["all"], user_id=None + ) + + assert auth.reserved == [] and auth.released == [] + assert indexing.added[0]["quota_reserved"] is False + + +@pytest.mark.asyncio +async def test_copy_file_reserves_and_consumes_the_slot(): + auth, indexing = FakeAuth(), FakeIndexing(copy_created=True) + svc = _service(auth=auth, indexing=indexing, partitions=_CopySourceOnly()) + + await svc.copy_file( + source_partition="p1", + source_file_id="src", + dest_partition="p1", + dest_file_id="dst", + allowed_partitions=["all"], + user_id=7, + ) + + assert auth.reserved == [7] and auth.released == [] + + +@pytest.mark.asyncio +async def test_copy_file_releases_the_slot_when_no_row_is_created(): + """An empty source writes no catalog row, so the reservation is unconsumed.""" + auth, indexing = FakeAuth(), FakeIndexing(copy_created=False) + svc = _service(auth=auth, indexing=indexing, partitions=_CopySourceOnly()) + + await svc.copy_file( + source_partition="p1", + source_file_id="src", + dest_partition="p1", + dest_file_id="dst", + allowed_partitions=["all"], + user_id=7, + ) + + assert auth.reserved == [7] and auth.released == [7] + + +@pytest.mark.asyncio +async def test_copy_file_over_quota_is_rejected_before_copying(): + auth, indexing = FakeAuth(over_quota=True), FakeIndexing() + svc = _service(auth=auth, indexing=indexing, partitions=_CopySourceOnly()) + + with pytest.raises(OpenRAGError) as exc: + await svc.copy_file( + source_partition="p1", + source_file_id="src", + dest_partition="p1", + dest_file_id="dst", + allowed_partitions=["all"], + user_id=7, + ) + + assert exc.value.code == "FILE_QUOTA_EXCEEDED" + assert indexing.copied == [] diff --git a/tests/unit/services/orchestrators/test_user_service.py b/tests/unit/services/orchestrators/test_user_service.py index dc80045e6..d8f5827dd 100644 --- a/tests/unit/services/orchestrators/test_user_service.py +++ b/tests/unit/services/orchestrators/test_user_service.py @@ -385,18 +385,35 @@ async def test_current_user_info_specific_quota_and_pending(): assert out["file_count"] == 4 assert out["pending_files"] == 3 - assert out["total_files"] == 7 + # #664: file_count already includes slots reserved at admission, so the + # in-flight uploads counted by ``pending_files`` are inside it. Adding + # them again would report 7/5 for a user the gate considers 4/5. + assert out["total_files"] == 4 assert out["file_quota"] == 5 assert out["id"] == 7 # original fields preserved assert job.calls == [7] +@pytest.mark.asyncio +async def test_current_user_info_pending_is_never_added_to_total(): + """Regression guard for the #664 double-count. + + ``pending_files`` is informational; the durable, already-reserved + ``file_count`` is the only number the quota gate acts on, so it must be + what ``total_files`` reports no matter how many tasks are in flight. + """ + svc = _svc(FakeUserRepo(), default_quota=10, job_service=FakeJobService(pending=99)) + out = await svc.get_current_user_info({"id": 7, "is_admin": False, "file_quota": 5, "file_count": 2}) + assert out["pending_files"] == 99 + assert out["total_files"] == 2 + + @pytest.mark.asyncio async def test_current_user_info_admin_is_unlimited(): svc = _svc(FakeUserRepo(), default_quota=10, job_service=FakeJobService(pending=1)) out = await svc.get_current_user_info({"id": 1, "is_admin": True, "file_count": 2}) assert out["file_quota"] == -1 - assert out["total_files"] == 3 + assert out["total_files"] == 2 @pytest.mark.asyncio diff --git a/tests/unit/services/persistence/test_job_repo.py b/tests/unit/services/persistence/test_job_repo.py new file mode 100644 index 000000000..b4877a94a --- /dev/null +++ b/tests/unit/services/persistence/test_job_repo.py @@ -0,0 +1,260 @@ +"""Unit tests for the durable :class:`PgJobRepository` (issue #660). + +The SQL is exercised against a fake asyncpg pool: these assert the shape of +the statements and the row -> model mapping. End-to-end behaviour against a +real Postgres lives in ``tests/integration/repos/``. +""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest +from core.models.catalog import DocumentStatus, IndexationJob + + +class _FakePool: + def __init__(self, row=None, rows=None, val=None): + self.executed: list[tuple[str, tuple]] = [] + self._row = row + self._rows = rows or [] + self._val = val + + async def fetchrow(self, query: str, *params): + self.executed.append((query, params)) + return self._row + + async def fetch(self, query: str, *params): + self.executed.append((query, params)) + return self._rows + + async def fetchval(self, query: str, *params): + self.executed.append((query, params)) + return self._val + + async def execute(self, query: str, *params): + self.executed.append((query, params)) + return "DELETE 3" + + +def _db_row(**overrides) -> dict: + row = { + "id": "task-1", + "status": "QUEUED", + "partition": "tenant-a", + "file_id": "file-1", + "user_id": 42, + "job_metadata": {"filename": "report.pdf"}, + "error": None, + "created_at": datetime(2026, 1, 1, tzinfo=UTC), + "started_at": None, + "completed_at": None, + "updated_at": datetime(2026, 1, 1, tzinfo=UTC), + } + row.update(overrides) + return row + + +def _repo(pool): + from services.persistence.job_repo import PgJobRepository + + return PgJobRepository(pool_getter=lambda: pool) + + +async def test_create_job_inserts_row_and_returns_model(): + pool = _FakePool(row=_db_row()) + job = await _repo(pool).create_job( + IndexationJob( + id="task-1", + status=DocumentStatus.QUEUED, + partition="tenant-a", + file_id="file-1", + user_id=42, + job_metadata={"filename": "report.pdf"}, + ) + ) + + query, params = pool.executed[0] + assert "INSERT INTO jobs" in query + assert params[0] == "task-1" + assert params[1] == "QUEUED" + assert isinstance(job, IndexationJob) + assert job.id == "task-1" + assert job.status is DocumentStatus.QUEUED + assert job.job_metadata == {"filename": "report.pdf"} + + +async def test_create_job_is_idempotent_on_redispatch(): + pool = _FakePool(row=_db_row()) + await _repo(pool).create_job(IndexationJob(id="task-1")) + + query, _ = pool.executed[0] + assert "ON CONFLICT (id) DO" in query + + +async def test_get_job_returns_none_when_missing(): + pool = _FakePool(row=None) + assert await _repo(pool).get_job("nope") is None + + +async def test_update_job_sets_only_allowlisted_fields(): + pool = _FakePool(row=_db_row(status="COMPLETED")) + job = await _repo(pool).update_job("task-1", status=DocumentStatus.COMPLETED, bogus="x") + + query, params = pool.executed[0] + assert "UPDATE jobs" in query + assert "bogus" not in query + assert "QUEUED" not in params + assert job.status is DocumentStatus.COMPLETED + + +async def test_update_job_upper_cases_a_lower_case_status(): + """A lower-case status must never reach the ``ck_jobs_status`` CHECK. + + ``update_job`` takes ``**fields`` from the worker path, so unlike + ``create_job`` (whose ``IndexationJob.status`` pydantic validates to a + ``DocumentStatus``) it can receive a plain string. Every durable write is + best-effort, so a CHECK violation here would be swallowed and the job would + silently freeze at its previous status — permanently, if the dropped write + was the terminal one. + """ + pool = _FakePool(row=_db_row(status="COMPLETED")) + await _repo(pool).update_job("task-1", status="completed") + + _, params = pool.executed[0] + assert "COMPLETED" in params + assert "completed" not in params + + +async def test_update_job_truncates_error_text(): + pool = _FakePool(row=_db_row(status="FAILED", error="x")) + await _repo(pool).update_job("task-1", status=DocumentStatus.FAILED, error="y" * 50_000) + + _, params = pool.executed[0] + stored = next(p for p in params if isinstance(p, str) and p.startswith("[")) + assert len(stored) < 10_000 + assert "truncated" in stored + + +async def test_mark_failed_if_not_cancelled_guards_on_status_in_sql(): + """The CANCELLED check must be part of the UPDATE, not a read-then-write. + + Arbitrating in the statement is what lets a worker whose state actor has + forgotten the task still reach a terminal row (#660): the actor can no + longer veto the write, and a concurrent cancel is still respected. + """ + pool = _FakePool(row={"id": "task-1"}) + assert await _repo(pool).mark_failed_if_not_cancelled( + "task-1", error="boom", completed_at=datetime(2026, 1, 1, tzinfo=UTC) + ) + + query, params = pool.executed[0] + assert "UPDATE jobs" in query + assert "status <> 'CANCELLED'" in query + assert "task-1" in params + + +async def test_mark_failed_if_not_cancelled_reports_a_cancelled_row(): + pool = _FakePool(row=None) # WHERE matched nothing: already CANCELLED, or gone + + assert not await _repo(pool).mark_failed_if_not_cancelled( + "task-1", error="boom", completed_at=datetime(2026, 1, 1, tzinfo=UTC) + ) + + +async def test_mark_failed_if_not_cancelled_truncates_error_text(): + pool = _FakePool(row={"id": "task-1"}) + await _repo(pool).mark_failed_if_not_cancelled( + "task-1", error="y" * 50_000, completed_at=datetime(2026, 1, 1, tzinfo=UTC) + ) + + _, params = pool.executed[0] + stored = next(p for p in params if isinstance(p, str) and p.startswith("[")) + assert len(stored) < 10_000 + assert "truncated" in stored + + +async def test_update_job_with_no_known_fields_is_a_noop_read(): + pool = _FakePool(row=_db_row()) + job = await _repo(pool).update_job("task-1", bogus="x") + + query, _ = pool.executed[0] + assert query.strip().startswith("SELECT") + assert job.id == "task-1" + + +async def test_list_jobs_filters_by_user_and_status(): + pool = _FakePool(rows=[_db_row()]) + jobs = await _repo(pool).list_jobs(status="FAILED", user_id=42, offset=5, limit=10) + + query, params = pool.executed[0] + assert "FROM jobs" in query + assert ["FAILED"] in params + assert 42 in params + assert 10 in params and 5 in params + assert len(jobs) == 1 + + +async def test_list_jobs_active_expands_to_the_non_terminal_states(): + pool = _FakePool(rows=[]) + await _repo(pool).list_jobs(status="active") + + _, params = pool.executed[0] + assert ["QUEUED", "SERIALIZING", "CHUNKING", "INSERTING"] in params + + +async def test_list_jobs_status_match_is_case_insensitive(): + pool = _FakePool(rows=[]) + await _repo(pool).list_jobs(status="failed") + + _, params = pool.executed[0] + assert ["FAILED"] in params + + +async def test_count_by_status_returns_mapping(): + pool = _FakePool(rows=[{"status": "COMPLETED", "count": 3}, {"status": "FAILED", "count": 1}]) + counts = await _repo(pool).count_by_status() + + assert counts == {"COMPLETED": 3, "FAILED": 1} + + +async def test_purge_terminal_jobs_deletes_aged_and_overflow_rows(): + pool = _FakePool(val=3) + purged = await _repo(pool).purge_terminal_jobs(older_than_seconds=3600, keep_last=100) + + assert purged == 3 + query, params = pool.executed[0] + assert "DELETE FROM jobs" in query + assert ["COMPLETED", "FAILED", "CANCELLED"] in params + assert 3600 in params + assert 100 in params + + +async def test_purge_terminal_jobs_rejects_negative_bounds(): + with pytest.raises(ValueError): + await _repo(_FakePool()).purge_terminal_jobs(older_than_seconds=-1, keep_last=10) + + +async def test_an_unwritable_status_is_rejected_before_it_reaches_sql(): + """A status outside the CHECK must fail loudly here, not silently in SQL. + + Every durable write is best-effort, so a ``CheckViolationError`` from + ``ck_jobs_status`` is swallowed by the caller and the job freezes at its + previous status — permanently, if the dropped write was the terminal one. + Casing is not the only way to build such a value: ``update_job`` is untyped, + and ``str(None).upper()`` is ``"NONE"``. + """ + pool = _FakePool(row=_db_row()) + for bad in (None, "bogus", "RUNNING"): + with pytest.raises(ValueError, match="unknown job status"): + await _repo(pool).update_job("task-1", status=bad) + + assert pool.executed == [], "a rejected status must not reach the database" + + +async def test_every_allowed_status_survives_the_guard(): + """The other half: the seven states the CHECK allows must all pass.""" + for state in DocumentStatus: + pool = _FakePool(row=_db_row(status=state.value)) + assert await _repo(pool).update_job("task-1", status=state) is not None + assert await _repo(pool).update_job("task-1", status=state.value.lower()) is not None diff --git a/tests/unit/services/workers/test_dispatcher.py b/tests/unit/services/workers/test_dispatcher.py index 6a09425d5..ed1cf183b 100644 --- a/tests/unit/services/workers/test_dispatcher.py +++ b/tests/unit/services/workers/test_dispatcher.py @@ -170,6 +170,7 @@ async def test_dispatch_indexing_queues_worker_pool_task_and_records_ref() -> No replace=True, indexation_config={"parsing_strategy": "pymupdf"}, embedder_name="embed-fast", + quota_reserved=False, require_existing_partition=True, ) tsm.set_object_ref.remote.assert_called_once_with("task-1", {"ref": ref}) @@ -1341,3 +1342,394 @@ async def test_delete_file_reports_failure_when_post_delete_cleanup_fails() -> N workspace_repo.remove_file_from_all_workspaces.assert_called_once_with("file-1", "tenant-a") document_repo.remove_file_from_partition.assert_called_once_with(file_id="file-1", partition="tenant-a") assert vector_store.delete_by_filter.await_count == 2 + + +# --------------------------------------------------------------------------- +# Durable job records (issue #660) +# --------------------------------------------------------------------------- + + +def _job_repo() -> MagicMock: + repo = MagicMock() + repo.create_job = AsyncMock(side_effect=lambda job: job) + repo.update_job = AsyncMock(return_value=None) + repo.get_job = AsyncMock(return_value=None) + repo.purge_terminal_jobs = AsyncMock(return_value=0) + return repo + + +def _dispatcher_with_job_repo(job_repo: Any, tsm: Any = None, ref: object | None = None) -> Any: + from services.workers.dispatcher import WorkerDispatcher + + return WorkerDispatcher( + pool=_pool_with_ref(ref if ref is not None else object()), + task_state_manager=tsm or _task_state_manager(), + vector_store=_vector_store(), + document_repo=_document_repo(), + workspace_repo=_workspace_repo(), + collection="default", + job_repo=job_repo, + ) + + +@pytest.mark.asyncio +async def test_dispatch_indexing_persists_a_queued_job_before_submitting() -> None: + from core.models.catalog import DocumentStatus + + job_repo = _job_repo() + dispatcher = _dispatcher_with_job_repo(job_repo) + + with patch("services.workers.dispatcher.uuid") as mock_uuid: + mock_uuid.uuid4.return_value.hex = "task-1" + await dispatcher.dispatch_indexing( + path="/data/report.txt", + metadata={"file_id": "file-1", "source": "/data/report.txt", "filename": "report.txt"}, + partition="tenant-a", + user={"id": 42}, + workspace_ids=None, + replace=False, + ) + + job = job_repo.create_job.await_args.args[0] + assert job.id == "task-1" + assert job.status is DocumentStatus.QUEUED + assert job.partition == "tenant-a" + assert job.file_id == "file-1" + assert job.user_id == 42 + assert job.job_metadata == {"filename": "report.txt"} + + +@pytest.mark.asyncio +async def test_dispatch_indexing_survives_a_job_repo_outage() -> None: + job_repo = _job_repo() + job_repo.create_job = AsyncMock(side_effect=RuntimeError("postgres down")) + dispatcher = _dispatcher_with_job_repo(job_repo) + + task_id = await dispatcher.dispatch_indexing( + path="/data/report.txt", + metadata={"file_id": "file-1"}, + partition="tenant-a", + user=None, + workspace_ids=None, + replace=False, + ) + + assert task_id + # indexing still went out to the pool despite the durable write failing + dispatcher._pool.submit.remote.assert_called_once() + + +@pytest.mark.asyncio +async def test_cancel_task_marks_the_durable_job_cancelled() -> None: + from core.models.catalog import DocumentStatus + + job_repo = _job_repo() + dispatcher = _dispatcher_with_job_repo(job_repo) + + with patch("ray.cancel"): + assert await dispatcher.cancel_task("task-1") is True + + job_repo.update_job.assert_awaited_once() + assert job_repo.update_job.await_args.args[0] == "task-1" + assert job_repo.update_job.await_args.kwargs["status"] is DocumentStatus.CANCELLED + assert job_repo.update_job.await_args.kwargs["completed_at"] is not None + + +@pytest.mark.asyncio +async def test_cancel_task_leaves_the_durable_job_alone_when_already_terminal() -> None: + """A cancel that loses the race must not rewrite a COMPLETED job as CANCELLED. + + ``get_object_ref`` still answers for a finished task (the ref is only dropped + when the entry is evicted), so a late ``DELETE /task/{id}`` reaches this path + for work that already succeeded. The durable row is the operator-visible + record of the outcome (#660), so it must reflect what actually happened. + """ + tsm = _task_state_manager() + tsm.set_cancelled_if_active = _remote_mock(False) # worker got there first + job_repo = _job_repo() + dispatcher = _dispatcher_with_job_repo(job_repo, tsm=tsm) + + with patch("ray.cancel") as cancel: + assert await dispatcher.cancel_task("task-1") is False + + job_repo.update_job.assert_not_awaited() + cancel.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_task_state_falls_back_to_postgres_after_a_restart() -> None: + from core.models.catalog import DocumentStatus, IndexationJob + + tsm = _task_state_manager() + tsm.get_state.remote = AsyncMock(return_value=None) # cache lost the entry + job_repo = _job_repo() + job_repo.get_job = AsyncMock(return_value=IndexationJob(id="task-1", status=DocumentStatus.COMPLETED)) + dispatcher = _dispatcher_with_job_repo(job_repo, tsm=tsm) + + assert await dispatcher.get_task_state("task-1") == "COMPLETED" + + +@pytest.mark.asyncio +async def test_get_task_error_falls_back_to_postgres_after_a_restart() -> None: + from core.models.catalog import IndexationJob + + tsm = _task_state_manager() + tsm.get_error.remote = AsyncMock(return_value=None) + job_repo = _job_repo() + job_repo.get_job = AsyncMock(return_value=IndexationJob(id="task-1", error="boom")) + dispatcher = _dispatcher_with_job_repo(job_repo, tsm=tsm) + + assert await dispatcher.get_task_error("task-1") == "boom" + + +@pytest.mark.asyncio +async def test_hot_cache_hit_does_not_query_postgres() -> None: + job_repo = _job_repo() + dispatcher = _dispatcher_with_job_repo(job_repo) + + assert await dispatcher.get_task_state("task-1") == "SERIALIZING" + job_repo.get_job.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_dispatch_purges_terminal_jobs_at_most_once_per_interval() -> None: + job_repo = _job_repo() + dispatcher = _dispatcher_with_job_repo(job_repo) + + for _ in range(3): + await dispatcher.dispatch_indexing( + path="/data/report.txt", + metadata={"file_id": "file-1"}, + partition="tenant-a", + user=None, + workspace_ids=None, + replace=False, + ) + + assert job_repo.purge_terminal_jobs.await_count == 1 + + +@pytest.mark.asyncio +async def test_the_purge_runs_again_once_the_interval_has_passed(monkeypatch) -> None: + """The throttle must rate-limit the sweep, not disable it after one run. + + Without this, an implementation that purges exactly once per process — and + so lets the table grow unbounded forever after — passes the at-most-once + test above. + """ + from services.workers import dispatcher as dispatcher_module + + clock = {"now": 1_000.0} + monkeypatch.setattr(dispatcher_module.time, "monotonic", lambda: clock["now"]) + job_repo = _job_repo() + dispatcher = _dispatcher_with_job_repo(job_repo) + + async def _dispatch(): + await dispatcher.dispatch_indexing( + path="/data/report.txt", + metadata={"file_id": "file-1"}, + partition="tenant-a", + user=None, + workspace_ids=None, + replace=False, + ) + + await _dispatch() + assert job_repo.purge_terminal_jobs.await_count == 1 + + clock["now"] += dispatcher_module.JOB_PURGE_INTERVAL_SECONDS - 1 + await _dispatch() + assert job_repo.purge_terminal_jobs.await_count == 1, "still inside the interval" + + clock["now"] += 2 + await _dispatch() + assert job_repo.purge_terminal_jobs.await_count == 2, "the interval has elapsed" + + +@pytest.mark.asyncio +async def test_the_purge_uses_the_documented_retention_bounds() -> None: + """Pin the shipped retention window and row cap. + + These are the only thing standing between the durable store and the + unbounded growth #660 exists to fix, and no other test asserts their values. + """ + job_repo = _job_repo() + dispatcher = _dispatcher_with_job_repo(job_repo) + + await dispatcher.dispatch_indexing( + path="/data/report.txt", + metadata={"file_id": "file-1"}, + partition="tenant-a", + user=None, + workspace_ids=None, + replace=False, + ) + + job_repo.purge_terminal_jobs.assert_awaited_once_with( + older_than_seconds=7 * 24 * 3600, + keep_last=10_000, + ) + + +@pytest.mark.asyncio +async def test_a_failing_purge_is_not_retried_on_every_dispatch() -> None: + """The throttle timestamp is stamped *before* the sweep, not after it. + + Stamped afterwards, a purge that raises never records an attempt, so every + subsequent upload pays another failing round-trip to a database that is + already unhealthy — turning a bounded 5-minute sweep into per-request load + at exactly the worst moment. ``test_purge_failure_never_fails_a_dispatch`` + covers that the failure is swallowed; this covers that it is not repeated. + """ + job_repo = _job_repo() + job_repo.purge_terminal_jobs = AsyncMock(side_effect=RuntimeError("purge blew up")) + dispatcher = _dispatcher_with_job_repo(job_repo) + + for _ in range(3): + await dispatcher.dispatch_indexing( + path="/data/report.txt", + metadata={"file_id": "file-1"}, + partition="tenant-a", + user=None, + workspace_ids=None, + replace=False, + ) + + assert job_repo.purge_terminal_jobs.await_count == 1, "a failing purge was retried on every dispatch" + + +@pytest.mark.asyncio +async def test_purge_failure_never_fails_a_dispatch() -> None: + job_repo = _job_repo() + job_repo.purge_terminal_jobs = AsyncMock(side_effect=RuntimeError("purge blew up")) + dispatcher = _dispatcher_with_job_repo(job_repo) + + assert await dispatcher.dispatch_indexing( + path="/data/report.txt", + metadata={"file_id": "file-1"}, + partition="tenant-a", + user=None, + workspace_ids=None, + replace=False, + ) + + +@pytest.mark.asyncio +async def test_cancel_writes_the_durable_row_before_killing_the_worker() -> None: + """The durable CANCELLED must be written before ``ray.cancel``, not after. + + ``ray.cancel`` kills the only other writer of the row, and the write that + follows it has no successor that could heal it: ``_record_job`` catches + ``Exception``, but a client disconnect raises ``asyncio.CancelledError`` — + a ``BaseException`` — straight through. Writing after the kill therefore + left the actor CANCELLED and the row stuck on its last active status + forever: non-terminal, so retention never sweeps it and it is counted + active for good, while the actor-first and durable-first read paths answer + differently for the same task id. + """ + order: list[str] = [] + job_repo = _job_repo() + job_repo.update_job = AsyncMock(side_effect=lambda *a, **k: order.append("durable")) + dispatcher = _dispatcher_with_job_repo(job_repo) + + with patch("ray.cancel", side_effect=lambda *a, **k: order.append("ray.cancel")): + assert await dispatcher.cancel_task("task-1") is True + + assert order == ["durable", "ray.cancel"], order + + +@pytest.mark.asyncio +async def test_a_cancellation_during_the_durable_write_still_kills_the_worker() -> None: + """A ``BaseException`` out of the durable write must not skip ``ray.cancel``. + + The user asked for a cancellation and the actor already claimed it; leaving + the worker running would contradict both records. + """ + job_repo = _job_repo() + job_repo.update_job = AsyncMock(side_effect=asyncio.CancelledError()) + dispatcher = _dispatcher_with_job_repo(job_repo) + + with patch("ray.cancel") as cancel: + with pytest.raises(asyncio.CancelledError): + await dispatcher.cancel_task("task-1") + + cancel.assert_called_once() + + +@pytest.mark.asyncio +async def test_a_failed_submit_settles_the_task_terminally_in_both_stores() -> None: + """A dispatch that never reaches a worker must not leave a QUEUED orphan. + + Nothing else can reclaim it: actor eviction is driven entirely off + ``terminal_at`` (which only a terminal state enters) and the detached actor + survives API restarts, while ``purge_terminal_jobs`` sweeps terminal rows + only. One permanent leak in each store per failed dispatch — verbatim the + unbounded growth #660 exists to fix. + """ + tsm = _task_state_manager() + tsm.set_failed_if_not_cancelled = _remote_mock(True) + job_repo = _job_repo() + job_repo.mark_failed_if_not_cancelled = AsyncMock(return_value=True) + dispatcher = _dispatcher_with_job_repo(job_repo, tsm=tsm) + dispatcher._pool.submit = _remote_mock() + dispatcher._pool.submit.remote = AsyncMock(side_effect=RuntimeError("pool is down")) + + with pytest.raises(RuntimeError, match="pool is down"): + await dispatcher.dispatch_indexing( + path="/tmp/f.txt", + metadata={"file_id": "f1"}, + partition="default", + user={"id": 7}, + workspace_ids=None, + replace=False, + ) + + tsm.set_failed_if_not_cancelled.remote.assert_awaited_once() + job_repo.mark_failed_if_not_cancelled.assert_awaited_once() + assert job_repo.mark_failed_if_not_cancelled.await_args.kwargs["completed_at"] is not None + + +@pytest.mark.asyncio +async def test_settling_a_failed_dispatch_never_masks_the_dispatch_error() -> None: + """The caller must see why the dispatch failed, not why the cleanup did.""" + tsm = _task_state_manager() + tsm.set_failed_if_not_cancelled = _remote_mock() + tsm.set_failed_if_not_cancelled.remote = AsyncMock(side_effect=RuntimeError("actor is gone")) + job_repo = _job_repo() + job_repo.mark_failed_if_not_cancelled = AsyncMock(side_effect=RuntimeError("postgres is down")) + dispatcher = _dispatcher_with_job_repo(job_repo, tsm=tsm) + dispatcher._pool.submit = _remote_mock() + dispatcher._pool.submit.remote = AsyncMock(side_effect=RuntimeError("pool is down")) + + with pytest.raises(RuntimeError, match="pool is down"): + await dispatcher.dispatch_indexing( + path="/tmp/f.txt", + metadata={"file_id": "f1"}, + partition="default", + user={"id": 7}, + workspace_ids=None, + replace=False, + ) + + +@pytest.mark.asyncio +async def test_an_empty_source_copies_nothing_and_reports_no_row(): + """A source with no chunks must report False so the slot goes back. + + ``copy_file``'s return value is what the router commits the reservation + on (#664): a True here would commit a slot for a copy that never wrote a + catalog row, permanently narrowing the user's quota. The MCP copy path + has this pinned; the dispatcher's own empty-source probe did not. + """ + dispatcher = _dispatcher_with_job_repo(_job_repo()) + dispatcher._vector_store.query_chunks_by_filter = AsyncMock(return_value=[]) + + created = await dispatcher.copy_file( + file_id="src", + metadata={}, + partition="p1", + user={"id": 42}, + ) + + assert created is False + dispatcher._document_repo.add_file_to_partition.assert_not_awaited() diff --git a/tests/unit/services/workers/test_indexer_pool.py b/tests/unit/services/workers/test_indexer_pool.py index 80a816dbd..c3537fb72 100644 --- a/tests/unit/services/workers/test_indexer_pool.py +++ b/tests/unit/services/workers/test_indexer_pool.py @@ -1011,6 +1011,9 @@ def model_copy(self, *, update): class Store: document_repo = object() topic_tag_repo = object() + # #664: the worker needs a user_repo to release reserved quota slots. + user_repo = object() + job_repo = object() class Worker: def __init__(self, **kwargs): @@ -1067,9 +1070,11 @@ class _RecordingWorker: def __init__(self, *, error: Exception | None = None) -> None: self._error = error self.calls = 0 + self.last_kwargs: dict = {} - async def process_file(self, **_kwargs) -> dict: + async def process_file(self, **kwargs) -> dict: self.calls += 1 + self.last_kwargs = kwargs if self._error is not None: raise self._error return {"stored_count": 1, "stage": "stored"} @@ -1088,7 +1093,8 @@ async def _noop(*_a, **_k): actor._ensure_catalog = _noop actor._ensure_registry_fresh = _noop actor._worker = worker - actor._catalog_store = SimpleNamespace(workspace_repo=SimpleNamespace()) + actor._catalog_store = SimpleNamespace(workspace_repo=SimpleNamespace(), job_repo=None) + actor._task_state_manager = _FakeStateManager() actor._save_uploaded_files = save_uploaded_files actor._logger = SimpleNamespace(debug=lambda *a, **k: None, warning=lambda *a, **k: None) return actor @@ -1170,3 +1176,211 @@ async def _boom(*_a, **_k): await actor.process_file(task_id="t", path=str(path), metadata={"file_id": "f"}, partition="p") assert path.exists() + + +# --------------------------------------------------------------------------- +# Setup-failure quota release (#664) +# --------------------------------------------------------------------------- + + +class _FakeUserRepo: + def __init__(self) -> None: + self.released: list[int] = [] + + async def release_file_slot(self, user_id: int) -> None: + self.released.append(user_id) + + +class _FakeRemote: + """Stands in for a Ray ``.remote`` handle method.""" + + def __init__(self, error: BaseException | None = None) -> None: + self.calls: list[tuple] = [] + self._error = error + + async def remote(self, *args): + self.calls.append(args) + if self._error is not None: + raise self._error + return True + + +class _FakeStateManager: + def __init__(self, error: BaseException | None = None) -> None: + self.set_failed_if_not_cancelled = _FakeRemote(error) + + +class _FakeJobRepo: + def __init__(self, error: BaseException | None = None) -> None: + self.failed: list[str] = [] + self._error = error + + async def mark_failed_if_not_cancelled(self, job_id, *, error, completed_at): + if self._error is not None: + raise self._error + self.failed.append(job_id) + return True + + +class _FakeCatalogStore: + def __init__(self, user_repo, job_repo=None) -> None: + self.user_repo = user_repo + self.job_repo = job_repo + + +def _pool_with_broken_setup(user_repo, *, error: BaseException, tsm=None, job_repo=None): + """A pool whose ``_ensure_catalog`` fails before the worker can take over.""" + from services.workers.indexer_pool import IndexerWorkerActor + + cls = IndexerWorkerActor.__ray_metadata__.modified_class + pool = cls.__new__(cls) + pool._catalog_store = _FakeCatalogStore(user_repo, job_repo) + pool._task_state_manager = tsm if tsm is not None else _FakeStateManager() + + async def _boom() -> None: + raise error + + async def _fresh(_names) -> None: + return None + + pool._ensure_catalog = _boom + pool._ensure_registry_fresh = _fresh + pool._worker = None # must never be reached + # Keep the upload on disk so the cleanup finally is a no-op here; the raw-file + # purge has its own coverage and is not what these tests are pinning. + pool._save_uploaded_files = True + pool._logger = None + return pool + + +@pytest.mark.parametrize( + "error", + [RuntimeError("catalog init failed"), asyncio.CancelledError()], + ids=["exception", "cancellation"], +) +async def test_setup_failure_releases_the_reserved_slot(error): + """Setup blew up before the worker owned the slot, so the pool must release it. + + The worker's own ``finally`` never runs when ``_ensure_catalog`` / + ``_ensure_registry_fresh`` raise, so without this release the upload + permanently narrows the user's quota. ``BaseException`` is deliberate: + a cancellation during setup must release too, and ``CancelledError`` is not + an ``Exception``. + """ + user_repo = _FakeUserRepo() + pool = _pool_with_broken_setup(user_repo, error=error) + + with pytest.raises(type(error)): + await pool.process_file( + task_id="t1", + path="/tmp/doc.txt", + metadata={"file_id": "f1"}, + partition="p1", + user={"id": 42}, + quota_reserved=True, + ) + + assert user_repo.released == [42], "a setup failure leaked the reserved quota slot" + + +async def test_setup_failure_releases_nothing_when_no_slot_was_reserved(): + """The other half: no reservation, no release (``put_file`` never reserves).""" + user_repo = _FakeUserRepo() + pool = _pool_with_broken_setup(user_repo, error=RuntimeError("catalog init failed")) + + with pytest.raises(RuntimeError): + await pool.process_file( + task_id="t1", + path="/tmp/doc.txt", + metadata={"file_id": "f1"}, + partition="p1", + user={"id": 42}, + quota_reserved=False, + ) + + assert user_repo.released == [] + + +@pytest.mark.asyncio +async def test_setup_failure_settles_the_task_terminally_in_both_stores(): + """A setup failure must not leave the task QUEUED forever. + + ``IndexerWorker`` -- the only writer of SERIALIZING/COMPLETED/FAILED -- is + never entered, so nothing else settles this task. Left QUEUED it is + unevictable in the detached actor (eviction reads ``terminal_at``, which + only a terminal state enters) and unsweepable by retention (terminal rows + only), so it is counted active in ``/queue/info`` for good -- while the + client polls a ``task_status_url`` that answers QUEUED forever. + """ + tsm = _FakeStateManager() + job_repo = _FakeJobRepo() + pool = _pool_with_broken_setup( + _FakeUserRepo(), error=RuntimeError("catalog init failed"), tsm=tsm, job_repo=job_repo + ) + + with pytest.raises(RuntimeError): + await pool.process_file( + task_id="t1", + path="/tmp/doc.txt", + metadata={"file_id": "f1"}, + partition="p1", + user={"id": 42}, + quota_reserved=True, + ) + + assert [c[0] for c in tsm.set_failed_if_not_cancelled.calls] == ["t1"] + assert job_repo.failed == ["t1"] + + +@pytest.mark.asyncio +async def test_settling_a_setup_failure_never_masks_the_setup_error(): + """The caller must see why setup failed, not why the bookkeeping did. + + Both settling writes go through stores that may be exactly what just + broke, so both have to be able to fail without changing the exception + the pool re-raises. + """ + tsm = _FakeStateManager(error=RuntimeError("state actor is gone")) + job_repo = _FakeJobRepo(error=RuntimeError("postgres is down")) + pool = _pool_with_broken_setup( + _FakeUserRepo(), error=RuntimeError("catalog init failed"), tsm=tsm, job_repo=job_repo + ) + + with pytest.raises(RuntimeError, match="catalog init failed"): + await pool.process_file( + task_id="t1", + path="/tmp/doc.txt", + metadata={"file_id": "f1"}, + partition="p1", + user={"id": 42}, + quota_reserved=True, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reserved", [True, False]) +async def test_the_actor_forwards_the_quota_reservation_to_the_worker(reserved): + """The pool must tell the worker whether it owns a reserved slot. + + ``quota_reserved`` is what arms the release in + ``IndexerWorker.process_file``'s ``finally`` (#664). The router end of + this wire is pinned by + ``test_dispatch_tells_the_worker_the_slot_is_already_reserved``, but the + pool hop was not: the setup-failure tests stub the worker out entirely + (``pool._worker = None # must never be reached``) and the worker tests + call ``process_file`` directly. Drop the keyword here and ``release_slot`` + is always False, so every failed upload leaks a slot in silence. + """ + worker = _RecordingWorker() + actor = _bare_worker_actor(save_uploaded_files=True, worker=worker) + + await actor.process_file( + task_id="t1", + path="/tmp/doc.txt", + metadata={"file_id": "f1"}, + partition="p1", + user={"id": 42}, + quota_reserved=reserved, + ) + + assert worker.last_kwargs["quota_reserved"] is reserved diff --git a/tests/unit/services/workers/test_indexer_worker.py b/tests/unit/services/workers/test_indexer_worker.py index a59b90936..b5608d9bf 100644 --- a/tests/unit/services/workers/test_indexer_worker.py +++ b/tests/unit/services/workers/test_indexer_worker.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from core.models.catalog import DocumentStatus from core.models.chunk import Chunk from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock from services.workers.indexer_actor import IndexerWorker, _load_document @@ -814,3 +815,244 @@ async def run(self, row: dict[str, Any]) -> dict[str, Any]: assert len(document_repo.add_calls) == 1 assert vector_store.deleted_filters == [] tsm.set_failed_if_not_cancelled.remote.assert_called_once() + + +# --------------------------------------------------------------------------- +# Durable job lifecycle writes (issue #660) +# --------------------------------------------------------------------------- + + +class FakeJobRepo: + def __init__(self) -> None: + self.updates: list[tuple[str, dict[str, Any]]] = [] + self.raise_on_update = False + # Job ids the durable store already holds as CANCELLED. Standing in for + # the ``status <> 'CANCELLED'`` guard that lives in SQL. + self.cancelled: set[str] = set() + + async def update_job(self, job_id: str, **fields: Any): + if self.raise_on_update: + raise RuntimeError("postgres down") + self.updates.append((job_id, fields)) + return None + + async def mark_failed_if_not_cancelled(self, job_id: str, *, error: str, completed_at: Any) -> bool: + if self.raise_on_update: + raise RuntimeError("postgres down") + if job_id in self.cancelled: + return False + self.updates.append((job_id, {"status": DocumentStatus.FAILED, "error": error, "completed_at": completed_at})) + return True + + +def _statuses(job_repo: FakeJobRepo) -> list[str]: + return [fields["status"].value for _, fields in job_repo.updates] + + +def _trivial_pipeline() -> Any: + return _make_pipeline( + ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="content")]), + [Chunk(id="c1", text="content", partition="p")], + ) + + +def _worker_with_job_repo(tmp_path: Path, job_repo: FakeJobRepo, pipeline: Any = None) -> IndexerWorker: + return IndexerWorker( + pipeline=pipeline if pipeline is not None else _trivial_pipeline(), + task_state_manager=_fake_tsm(), + job_repo=job_repo, + ) + + +@pytest.mark.asyncio +async def test_process_file_records_serializing_then_completed_durably(tmp_path: Path) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"hello") + job_repo = FakeJobRepo() + worker = _worker_with_job_repo(tmp_path, job_repo) + + await worker.process_file(task_id="t1", path=str(path), metadata={"file_id": "f1"}, partition="p") + + assert _statuses(job_repo) == ["SERIALIZING", "COMPLETED"] + assert job_repo.updates[0][1]["started_at"] is not None + assert job_repo.updates[-1][1]["completed_at"] is not None + + +@pytest.mark.asyncio +async def test_process_file_records_the_failure_and_its_traceback(tmp_path: Path) -> None: + """The traceback is handed over whole; the bound is applied by the store. + + ``PgJobRepository.update_job`` truncates ``error`` (see + ``tests/unit/services/persistence/test_job_repo.py``) so the cap holds for + every writer, not just this one. + """ + path = tmp_path / "doc.txt" + path.write_bytes(b"hello") + job_repo = FakeJobRepo() + + class BoomPipeline: + async def run(self, row): + raise RuntimeError("kaboom") + + worker = _worker_with_job_repo(tmp_path, job_repo, pipeline=BoomPipeline()) + + with pytest.raises(RuntimeError): + await worker.process_file(task_id="t1", path=str(path), metadata={"file_id": "f1"}, partition="p") + + assert _statuses(job_repo) == ["SERIALIZING", "FAILED"] + assert "kaboom" in job_repo.updates[-1][1]["error"] + assert job_repo.updates[-1][1]["completed_at"] is not None + + +@pytest.mark.asyncio +async def test_process_file_does_not_overwrite_a_cancelled_job(tmp_path: Path) -> None: + """A cancellation the user already saw survives a late failure write.""" + path = tmp_path / "doc.txt" + path.write_bytes(b"hello") + job_repo = FakeJobRepo() + job_repo.cancelled.add("t1") # the user cancelled while the pipeline ran + + class BoomPipeline: + async def run(self, row): + raise RuntimeError("boom") + + worker = IndexerWorker(pipeline=BoomPipeline(), task_state_manager=_fake_tsm(), job_repo=job_repo) + + with pytest.raises(RuntimeError): + await worker.process_file(task_id="t1", path=str(path), metadata={"file_id": "f1"}, partition="p") + + assert _statuses(job_repo) == ["SERIALIZING"] + + +@pytest.mark.asyncio +async def test_failure_is_recorded_even_when_the_state_actor_lost_the_task(tmp_path: Path) -> None: + """A forgetful state actor must not strand the durable row in SERIALIZING. + + ``set_failed_if_not_cancelled`` answers False both for a real cancellation + and for a task the actor no longer has (restart, TTL eviction). Gating the + durable write on it left the row non-terminal forever, since retention only + sweeps terminal rows. Postgres now arbitrates, so this still reaches FAILED. + """ + path = tmp_path / "doc.txt" + path.write_bytes(b"hello") + job_repo = FakeJobRepo() # nothing cancelled: the row is live + tsm = _fake_tsm() + tsm.set_failed_if_not_cancelled.remote = AsyncMock(return_value=False) # entry evicted + + class BoomPipeline: + async def run(self, row): + raise RuntimeError("boom") + + worker = IndexerWorker(pipeline=BoomPipeline(), task_state_manager=tsm, job_repo=job_repo) + + with pytest.raises(RuntimeError): + await worker.process_file(task_id="t1", path=str(path), metadata={"file_id": "f1"}, partition="p") + + assert _statuses(job_repo) == ["SERIALIZING", "FAILED"] + assert "boom" in job_repo.updates[-1][1]["error"] + + +@pytest.mark.asyncio +async def test_durable_write_failure_never_fails_indexing(tmp_path: Path) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"hello") + job_repo = FakeJobRepo() + job_repo.raise_on_update = True + worker = _worker_with_job_repo(tmp_path, job_repo) + + result = await worker.process_file(task_id="t1", path=str(path), metadata={"file_id": "f1"}, partition="p") + + assert result["stored_count"] == 1 + + +@pytest.mark.asyncio +async def test_an_unreachable_state_actor_still_writes_the_durable_failure(tmp_path: Path) -> None: + """The hot-cache write must not be able to take ``_mark_job_failed`` with it. + + The failure handler talks to the detached ``TaskStateManager`` first. Left + unguarded, an actor that is gone (restart, lost node) raises straight out of + the ``except`` block and the durable write never runs — leaving the row + non-terminal, which ``purge_terminal_jobs`` never sweeps. Every file + dispatched during an actor outage would strand one row forever: the + unbounded growth #660 exists to fix, on the durable side. + """ + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="content")]) + chunks = [Chunk(id="c1", text="content", partition="p")] + + tsm = _fake_tsm() + tsm.set_failed_if_not_cancelled.remote = AsyncMock(side_effect=RuntimeError("actor gone")) + + class BrokenRepo: + async def add_file_to_partition(self, **kwargs: Any) -> bool: + raise RuntimeError("pg down") + + class RecordingJobRepo: + def __init__(self) -> None: + self.failed: list[dict[str, Any]] = [] + + async def update_job(self, job_id: str, **fields: Any) -> None: + return None + + async def mark_failed_if_not_cancelled(self, job_id: str, **kwargs: Any) -> bool: + self.failed.append({"job_id": job_id, **kwargs}) + return True + + job_repo = RecordingJobRepo() + worker = IndexerWorker( + pipeline=_make_pipeline(processed, chunks), + task_state_manager=tsm, + document_repo=BrokenRepo(), + job_repo=job_repo, + ) + + # The original pipeline error must survive, not the actor's. + with pytest.raises(RuntimeError, match="pg down"): + await worker.process_file( + task_id="t-actor-gone", + path=str(path), + metadata={"file_id": "f1"}, + partition="p", + ) + + assert [row["job_id"] for row in job_repo.failed] == ["t-actor-gone"], ( + "the durable FAILED write was lost when the state actor was unreachable" + ) + + +@pytest.mark.asyncio +async def test_a_replace_reindex_updates_the_row_without_creating_one(): + """A ``replace`` re-index reuses an existing row, and must say the write landed. + + This pins the contract that the rebase onto #671 had to settle. Before it, + ``_write_catalog_record`` returned "a *new* row was created", so ``replace`` + reported False and the caller read that as "the reservation was not + consumed". #671 added a fail-closed ``if not wrote_catalog: raise`` on the + same value, which would have turned every ``replace`` re-index into a hard + failure. + + The two questions are now separate: this returns whether the catalog write + landed (True on either branch), and the quota verdict is derived from + ``replace`` at the call site. So the assertion here is that a replace + updates rather than inserts, and still reports success -- + ``test_replace_reindex_never_releases`` covers the quota half. + """ + from services.workers.indexer_actor import _write_catalog_record + + doc_repo = MagicMock() + doc_repo.update_file_in_partition = AsyncMock(return_value=True) + doc_repo.add_file_to_partition = AsyncMock(return_value=True) + + wrote_catalog = await _write_catalog_record( + doc_repo=doc_repo, + metadata={"file_id": "f1"}, + partition="p1", + user={"id": 42}, + replace=True, + indexation_config=None, + ) + + assert wrote_catalog is True, "a replace must report the write as landed, or #671 fails it closed" + doc_repo.update_file_in_partition.assert_awaited_once() + doc_repo.add_file_to_partition.assert_not_awaited() diff --git a/tests/unit/services/workers/test_indexer_worker_quota_release.py b/tests/unit/services/workers/test_indexer_worker_quota_release.py new file mode 100644 index 000000000..37f9ce60d --- /dev/null +++ b/tests/unit/services/workers/test_indexer_worker_quota_release.py @@ -0,0 +1,242 @@ +"""Issue #664 — the worker owns the reserved quota slot after dispatch. + +Admission charges one ``users.file_count`` slot *before* the job is queued, +so from dispatch onward the worker must either consume that slot (write a +catalog row) or give it back. A missed release leaks ``file_count`` upward +and permanently narrows the user's quota, so every non-success outcome gets +its own test here. +""" + +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from services.workers.indexer_actor import IndexerWorker + + +class FakeUserRepo: + def __init__(self, *, boom: bool = False) -> None: + self.released: list[int] = [] + self.boom = boom + + async def release_file_slot(self, user_id: int) -> None: + if self.boom: + raise RuntimeError("database is down") + self.released.append(user_id) + + +class FakeDocumentRepo: + def __init__(self, *, add_result: bool = True) -> None: + self.add_result = add_result + self.add_calls: list[dict[str, Any]] = [] + self.update_calls: list[dict[str, Any]] = [] + + async def add_file_to_partition(self, **kwargs: Any) -> bool: + self.add_calls.append(kwargs) + return self.add_result + + async def update_file_in_partition(self, **kwargs: Any) -> bool: + self.update_calls.append(kwargs) + return True + + +class _Pipeline: + """Minimal pipeline stand-in: succeed, raise, or hang until cancelled.""" + + def __init__(self, *, error: BaseException | None = None, hang: bool = False) -> None: + self.error = error + self.hang = hang + + async def run(self, row: dict) -> dict: + if self.hang: + await asyncio.sleep(3600) + if self.error is not None: + raise self.error + return {**row, "stored_count": 3, "stage": "stored", "indexed_at": None} + + +def _tsm() -> MagicMock: + tsm = MagicMock() + tsm.set_state = MagicMock() + tsm.set_state.remote = AsyncMock(return_value=None) + tsm.set_failed_if_not_cancelled = MagicMock() + tsm.set_failed_if_not_cancelled.remote = AsyncMock(return_value=True) + return tsm + + +def _worker(pipeline, doc_repo, user_repo) -> IndexerWorker: + return IndexerWorker( + pipeline=pipeline, + task_state_manager=_tsm(), + document_repo=doc_repo, + topic_tag_repo=None, + user_repo=user_repo, + ) + + +async def _run(worker, tmp_path, **overrides): + path = tmp_path / "doc.txt" + path.write_text("hello") + kwargs: dict[str, Any] = { + "task_id": "t1", + "path": str(path), + "metadata": {"file_id": "f1", "filename": "doc.txt"}, + "partition": "p1", + "user": {"id": 42}, + "quota_reserved": True, + } + kwargs.update(overrides) + return await worker.process_file(**kwargs) + + +@pytest.mark.asyncio +async def test_success_consumes_the_slot(tmp_path): + """The catalog row *is* the reservation — nothing to release.""" + user_repo = FakeUserRepo() + doc_repo = FakeDocumentRepo(add_result=True) + + result = await _run(_worker(_Pipeline(), doc_repo, user_repo), tmp_path) + + assert result["stored_count"] == 3 + assert len(doc_repo.add_calls) == 1 + assert user_repo.released == [] + + +@pytest.mark.asyncio +async def test_indexing_failure_releases_the_slot(tmp_path): + user_repo = FakeUserRepo() + doc_repo = FakeDocumentRepo() + + with pytest.raises(RuntimeError, match="embedder exploded"): + await _run(_worker(_Pipeline(error=RuntimeError("embedder exploded")), doc_repo, user_repo), tmp_path) + + assert user_repo.released == [42] + assert doc_repo.add_calls == [] + + +@pytest.mark.asyncio +async def test_cancellation_releases_the_slot(tmp_path): + """``ray.cancel`` raises CancelledError — a BaseException. + + Regression guard: the release must not live in an ``except Exception`` + block, which would sail straight past a cancelled task and leak the slot. + """ + user_repo = FakeUserRepo() + worker = _worker(_Pipeline(hang=True), FakeDocumentRepo(), user_repo) + + task = asyncio.create_task(_run(worker, tmp_path)) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert user_repo.released == [42] + + +@pytest.mark.asyncio +async def test_duplicate_at_catalog_releases_the_slot(tmp_path): + """The 409 check is pre-dispatch, so two racers can both reach the insert. + + ``add_file_to_partition`` returns False for the loser: no file was created + for its reservation, so the slot must go back. + + Since the rebase onto #671 the loser gets there by a different route. That + branch added a fail-closed ``if not wrote_catalog: raise`` on the catalog + write, so a False no longer falls through to a ``created``-gated release -- + it raises, and ``process_file``'s ``finally`` releases on the way out. The + invariant this test exists for is unchanged and still the point: whatever + the route, a reservation that produced no file row goes back. Both halves + are asserted so neither can regress silently. + """ + user_repo = FakeUserRepo() + doc_repo = FakeDocumentRepo(add_result=False) + + with pytest.raises(RuntimeError, match="Catalog row was not written"): + await _run(_worker(_Pipeline(), doc_repo, user_repo), tmp_path) + + assert len(doc_repo.add_calls) == 1 + assert user_repo.released == [42] + + +@pytest.mark.asyncio +async def test_replace_reindex_never_releases(tmp_path): + """``put_file`` reuses an existing row, so it never reserved a slot.""" + user_repo = FakeUserRepo() + doc_repo = FakeDocumentRepo() + + await _run(_worker(_Pipeline(), doc_repo, user_repo), tmp_path, replace=True, quota_reserved=False) + + assert doc_repo.update_calls and doc_repo.add_calls == [] + assert user_repo.released == [] + + +@pytest.mark.asyncio +async def test_failure_without_a_reservation_releases_nothing(tmp_path): + """A job dispatched without reserving must not decrement anyone.""" + user_repo = FakeUserRepo() + + with pytest.raises(RuntimeError): + await _run( + _worker(_Pipeline(error=RuntimeError("boom")), FakeDocumentRepo(), user_repo), + tmp_path, + quota_reserved=False, + ) + + assert user_repo.released == [] + + +@pytest.mark.asyncio +async def test_anonymous_upload_releases_nothing(tmp_path): + """No user id ⇒ nothing was charged ⇒ nothing to give back.""" + user_repo = FakeUserRepo() + + with pytest.raises(RuntimeError): + await _run( + _worker(_Pipeline(error=RuntimeError("boom")), FakeDocumentRepo(), user_repo), + tmp_path, + user=None, + ) + + assert user_repo.released == [] + + +@pytest.mark.asyncio +async def test_release_failure_does_not_mask_the_indexing_error(tmp_path): + """A broken release must not replace the real cause in the traceback.""" + user_repo = FakeUserRepo(boom=True) + + with pytest.raises(RuntimeError, match="embedder exploded"): + await _run( + _worker(_Pipeline(error=RuntimeError("embedder exploded")), FakeDocumentRepo(), user_repo), + tmp_path, + ) + + +@pytest.mark.asyncio +async def test_state_write_failure_releases_the_slot(tmp_path): + """The SERIALIZING write can fail — the slot must still go back. + + ``set_state`` talks to a detached Ray actor that can be unreachable + (restart, OOM, node loss). It runs *before* the try block that owns the + release, and ``IndexerWorkerActor.process_file`` only guards the + catalog/registry setup, so nothing else covers this window: the request + already handed the slot off at dispatch and will not release it. + """ + user_repo = FakeUserRepo() + tsm = _tsm() + tsm.set_state.remote = AsyncMock(side_effect=RuntimeError("actor unreachable")) + worker = IndexerWorker( + pipeline=_Pipeline(), + task_state_manager=tsm, + document_repo=FakeDocumentRepo(), + topic_tag_repo=None, + user_repo=user_repo, + ) + + with pytest.raises(RuntimeError, match="actor unreachable"): + await _run(worker, tmp_path) + + assert user_repo.released == [42] diff --git a/tests/unit/services/workers/test_task_state.py b/tests/unit/services/workers/test_task_state.py index 65b4e4075..0a9964e7c 100644 --- a/tests/unit/services/workers/test_task_state.py +++ b/tests/unit/services/workers/test_task_state.py @@ -3,6 +3,7 @@ from typing import Any import pytest +from services.workers import task_state as task_state_module from services.workers.task_state import PENDING_TASK_DETAILS, TaskStateManager @@ -12,6 +13,7 @@ def _task_state_manager() -> Any: @pytest.mark.asyncio async def test_cancelled_state_is_not_overwritten_by_worker_transitions() -> None: + """A cancel claim is sticky: a worker still in flight cannot report over it (#685).""" manager = _task_state_manager() await manager.set_state("task-1", "QUEUED") @@ -182,3 +184,247 @@ async def test_file_delete_fence_is_counted_for_overlapping_deletes() -> None: ) is True ) + + +_ACTOR = task_state_module.TaskStateManager.__ray_metadata__.modified_class + + +def _manager(): + return _ACTOR() + + +async def _dispatch(mgr, task_id: str, user_id: int = 1) -> None: + await mgr.set_state(task_id, "QUEUED") + await mgr.set_details(task_id, file_id=f"f-{task_id}", partition="p", metadata={}, user_id=user_id) + + +async def test_terminal_tasks_are_evicted_beyond_the_cap(monkeypatch): + monkeypatch.setattr(task_state_module, "_MAX_TERMINAL_TASKS", 3) + mgr = _manager() + + for i in range(6): + await _dispatch(mgr, f"t{i}") + await mgr.set_state(f"t{i}", "COMPLETED") + + assert len(mgr.tasks) == 3 + # oldest evicted first (FIFO) + assert set(mgr.tasks) == {"t3", "t4", "t5"} + assert await mgr.get_state("t0") is None + + +async def test_eviction_drops_the_user_index_entry_too(monkeypatch): + monkeypatch.setattr(task_state_module, "_MAX_TERMINAL_TASKS", 1) + mgr = _manager() + + await _dispatch(mgr, "t0", user_id=7) + await mgr.set_state("t0", "COMPLETED") + await _dispatch(mgr, "t1", user_id=7) + await mgr.set_state("t1", "COMPLETED") + + assert list(mgr.tasks) == ["t1"] + assert mgr.user_index.get(7) == {"t1"} + + +async def test_terminal_tasks_are_evicted_once_older_than_the_ttl(monkeypatch): + clock = {"now": 0.0} + monkeypatch.setattr(task_state_module.time, "monotonic", lambda: clock["now"]) + monkeypatch.setattr(task_state_module, "_TERMINAL_TTL_SECONDS", 10.0) + mgr = _manager() + + await _dispatch(mgr, "old") + await mgr.set_state("old", "COMPLETED") + + clock["now"] = 100.0 + await _dispatch(mgr, "new") + await mgr.set_state("new", "FAILED") + + assert "old" not in mgr.tasks + assert "new" in mgr.tasks + + +async def test_in_flight_tasks_are_never_evicted(monkeypatch): + monkeypatch.setattr(task_state_module, "_MAX_TERMINAL_TASKS", 1) + mgr = _manager() + + await _dispatch(mgr, "running") + await mgr.set_state("running", "SERIALIZING") + for i in range(5): + await _dispatch(mgr, f"done{i}") + await mgr.set_state(f"done{i}", "COMPLETED") + + assert await mgr.get_state("running") == "SERIALIZING" + + +async def test_cancelled_tasks_are_evictable(monkeypatch): + monkeypatch.setattr(task_state_module, "_MAX_TERMINAL_TASKS", 0) + mgr = _manager() + + await _dispatch(mgr, "t0") + await mgr.set_state("t0", "CANCELLED") + + assert mgr.tasks == {} + + +async def test_a_claimed_cancellation_is_evictable(monkeypatch): + """The *real* cancel entry point must register the terminal transition. + + ``test_cancelled_tasks_are_evictable`` above goes through ``set_state``, + which no cancellation actually uses: ``WorkerDispatcher.cancel_task`` claims + the cancellation with ``set_cancelled_if_active``, and nothing writes that + task's state again -- ``ray.cancel`` raises ``CancelledError``, a + ``BaseException`` that ``process_file``'s ``except Exception`` sails past. + If this path skips ``_mark_terminal`` the entry never enters + ``terminal_at``, so neither the cap nor the TTL can reclaim it, and every + user-initiated cancel leaks one ``TaskInfo`` -- with its details, its + ``user_index`` entry, and its pinned ``object_ref`` -- for the lifetime of + the detached actor. + """ + monkeypatch.setattr(task_state_module, "_MAX_TERMINAL_TASKS", 0) + mgr = _manager() + + await _dispatch(mgr, "t0") + assert await mgr.set_cancelled_if_active("t0") is True + + assert mgr.tasks == {} + assert mgr.terminal_at == {} + assert mgr.user_index == {} + + +async def test_set_error_truncates_long_tracebacks(): + mgr = _manager() + await _dispatch(mgr, "t0") + + await mgr.set_error("t0", "x" * 100_000) + + stored = await mgr.get_error("t0") + assert len(stored) <= task_state_module._MAX_ERROR_CHARS + 100 + assert "truncated" in stored + + +async def test_set_failed_if_not_cancelled_truncates_and_reports_terminality(): + mgr = _manager() + await _dispatch(mgr, "t0") + + assert await mgr.set_failed_if_not_cancelled("t0", "y" * 100_000) is True + assert await mgr.get_state("t0") == "FAILED" + assert len(await mgr.get_error("t0")) <= task_state_module._MAX_ERROR_CHARS + 100 + + +async def test_set_failed_if_not_cancelled_keeps_cancelled_state(): + mgr = _manager() + await _dispatch(mgr, "t0") + await mgr.set_state("t0", "CANCELLED") + + assert await mgr.set_failed_if_not_cancelled("t0", "boom") is False + assert await mgr.get_state("t0") == "CANCELLED" + + +@pytest.mark.parametrize("state", ["COMPLETED", "FAILED", "CANCELLED"]) +async def test_failing_after_eviction_does_not_resurrect_the_task(monkeypatch, state): + monkeypatch.setattr(task_state_module, "_MAX_TERMINAL_TASKS", 0) + mgr = _manager() + + await _dispatch(mgr, "t0") + await mgr.set_state("t0", state) + + assert await mgr.set_failed_if_not_cancelled("t0", "late") is False + assert mgr.tasks == {} + + +def _late_writes(): + """Every setter except ``set_state``, which alone may create a task.""" + return { + "set_error": lambda mgr: mgr.set_error("t0", "late"), + "set_details": lambda mgr: mgr.set_details("t0", file_id="f", partition="p", metadata={}, user_id=7), + "set_object_ref": lambda mgr: mgr.set_object_ref("t0", {"ref": object()}), + } + + +@pytest.mark.parametrize("writer", list(_late_writes())) +async def test_a_late_write_does_not_resurrect_an_evicted_task(monkeypatch, writer): + """A write for an evicted task is dropped, not turned into a new entry. + + Regression: these setters used to go through ``_ensure_task``, which + recreates on miss. The recreated ``TaskInfo`` has ``state=None``, so it never + enters ``terminal_at`` and is never evictable again — an unbounded leak on a + detached actor, which is the whole failure #660 exists to fix. + """ + monkeypatch.setattr(task_state_module, "_MAX_TERMINAL_TASKS", 0) + mgr = _manager() + + await _dispatch(mgr, "t0", user_id=7) + await mgr.set_state("t0", "COMPLETED") + assert mgr.tasks == {}, "precondition: the task is evicted" + + await _late_writes()[writer](mgr) + + assert mgr.tasks == {}, f"{writer} resurrected an evicted task" + assert mgr.terminal_at == {} + assert mgr.user_index == {} + + +async def test_a_late_write_for_an_unknown_task_is_dropped(): + """The same guard, without eviction: an id we never dispatched stays unknown.""" + mgr = _manager() + + await mgr.set_error("never-dispatched", "boom") + + assert mgr.tasks == {} + assert await mgr.get_state("never-dispatched") is None + + +async def test_the_shipped_bounds_are_the_documented_ones(): + """Pin the production values; every other test here monkeypatches them. + + Without this, narrowing the cap to something that no longer bounds memory — + or widening it back to the unbounded behaviour of #660 — breaks no test. + """ + assert task_state_module._MAX_TERMINAL_TASKS == 2000 + assert task_state_module._TERMINAL_TTL_SECONDS == 3600.0 + + +async def test_an_idle_deployment_keeps_its_last_terminal_task_cached(monkeypatch): + """The TTL is swept lazily, on settle — not on a timer. Reads do not check age. + + This pins the documented trade-off rather than an aspiration: with no new + task settling, an expired entry stays cached. That is harmless (a terminal + state is immutable, so a stale read is not a wrong read) and the cap still + bounds memory — but a reader must not assume the TTL has retired anything. + """ + clock = {"now": 0.0} + monkeypatch.setattr(task_state_module.time, "monotonic", lambda: clock["now"]) + monkeypatch.setattr(task_state_module, "_TERMINAL_TTL_SECONDS", 10.0) + mgr = _manager() + + await _dispatch(mgr, "only") + await mgr.set_state("only", "COMPLETED") + + clock["now"] = 10_000.0 # long past the TTL, but nothing else settles + + assert await mgr.get_state("only") == "COMPLETED" + assert "only" in mgr.tasks + + +async def test_a_failure_is_evictable(monkeypatch): + """The *primary* failure path must register the terminal transition too. + + ``test_a_claimed_cancellation_is_evictable`` pins this for the cancel + entry point, but the argument applies harder here: + ``IndexerWorker.process_file``'s ``except`` reaches for + ``set_failed_if_not_cancelled``, never ``set_state``, so this is the path + every failed indexing job actually takes. Skipping ``_mark_terminal`` + keeps the entry out of ``terminal_at``, so neither the cap nor the TTL + can reclaim it and every failure leaks one ``TaskInfo`` -- with its + stored traceback, its ``user_index`` entry and its pinned + ``object_ref`` -- for the lifetime of the detached actor. That is + verbatim the unbounded growth #660 exists to fix. + """ + monkeypatch.setattr(task_state_module, "_MAX_TERMINAL_TASKS", 0) + mgr = _manager() + + await _dispatch(mgr, "t0") + assert await mgr.set_failed_if_not_cancelled("t0", "boom") is True + + assert mgr.tasks == {} + assert mgr.terminal_at == {} + assert mgr.user_index == {}