diff --git a/config/config.default.yaml b/config/config.default.yaml index a09c6191..a046040e 100644 --- a/config/config.default.yaml +++ b/config/config.default.yaml @@ -63,6 +63,8 @@ mongo: authsource: attempts: 3 # Number of attempts to accomplish operation before considering it failed. shutdown_timeout: 3 # Time in seconds to wait for serialize task to finish. + compression_level: 10 # zstd level used to compress stored job documents. Higher means greater compression, at greater CPU cost. + max_stored_doc_bytes: 1073741824 # Compressed job-docs at or above this size are not persisted. tier0: backend: gandalf backend_infores: infores:dogpark-tier0 diff --git a/src/retriever/config/general.py b/src/retriever/config/general.py index 099c889f..3dad6956 100644 --- a/src/retriever/config/general.py +++ b/src/retriever/config/general.py @@ -93,6 +93,18 @@ class MongoSettings(BaseModel): shutdown_timeout: Annotated[ int, Field(description="Time in seconds to wait for serialize task to finish.") ] = 3 + compression_level: Annotated[ + int, + Field( + description="zstd level used to compress stored job documents. Higher means greater compression, at greater CPU cost." + ), + ] = 10 + max_stored_doc_bytes: Annotated[ + int, + Field( + description="Compressed job-docs at or above this size are not persisted." + ), + ] = 1024**3 class TelemetrySettings(BaseModel): diff --git a/src/retriever/lookup/lookup.py b/src/retriever/lookup/lookup.py index b6c6a5d8..cdf14994 100644 --- a/src/retriever/lookup/lookup.py +++ b/src/retriever/lookup/lookup.py @@ -47,7 +47,7 @@ MONGO_QUEUE = MongoQueue() OP_TABLE_MANAGER = OpTableManager() -ZSTD_COMPRESSOR = zstandard.ZstdCompressor() +ZSTD_COMPRESSOR = zstandard.ZstdCompressor(level=CONFIG.mongo.compression_level) CALLBACK_BODY_LOG_LIMIT = 500 """Truncate logged callback response bodies to this many chars.""" @@ -172,6 +172,10 @@ async def lookup(query: QueryInfo) -> tuple[HTTPStatus, ResponseDict]: job_log.info( f"Begin processing job {job_id} for client {get_submitter(query)}." ) + if (query.body.get("parameters") or {}).get("dehydrated"): + job_log.info( + "Query running in dehydrated mode. Response will not be stored." + ) qgraph = query.body["message"].get("query_graph") if qgraph is None: @@ -398,21 +402,22 @@ def tracked_response( # Cast because TypedDict has some *really annoying* interactions with more general dicts kgraph = response["message"].get("knowledge_graph") or {} + # Don't store dehydrated job responses, they're usually huge + dehydrated = bool(((query.body or {}).get("parameters") or {}).get("dehydrated")) + state = ResponseState( + job_id=query.job_id, + event_time=datetime.now().astimezone(), + knodes=len(kgraph.get("nodes") or {}), + kedges=len(kgraph.get("edges") or {}), + aux_graphs=len(response["message"].get("auxiliary_graphs") or {}), + results=len(response["message"].get("results") or []), + status=(response.get("status") or "Running"), + description=response.get("description"), + ) + if not dehydrated: + state["response"] = ZSTD_COMPRESSOR.compress(ormsgpack.packb(response)) try: - MONGO_QUEUE.put( - "job_state", - ResponseState( - job_id=query.job_id, - response=ZSTD_COMPRESSOR.compress(ormsgpack.packb(response)), - event_time=datetime.now().astimezone(), - knodes=len(kgraph.get("nodes") or {}), - kedges=len(kgraph.get("edges") or {}), - aux_graphs=len(response["message"].get("auxiliary_graphs") or {}), - results=len(response["message"].get("results") or []), - status=(response.get("status") or "Running"), - description=response.get("description"), - ), - ) + MONGO_QUEUE.put("job_state", state) except MongoOutage: response["logs"] = [ *(response.get("logs") or []), diff --git a/src/retriever/query.py b/src/retriever/query.py index 383da554..7955fea7 100644 --- a/src/retriever/query.py +++ b/src/retriever/query.py @@ -49,7 +49,7 @@ tracer = trace.get_tracer("lookup.execution.tracer") MONGO_QUEUE = MongoQueue() -ZSTD_COMPRESSOR = zstandard.ZstdCompressor() +ZSTD_COMPRESSOR = zstandard.ZstdCompressor(level=CONFIG.mongo.compression_level) ZSTD_DECOMPRESSOR = zstandard.ZstdDecompressor() @@ -119,6 +119,9 @@ def _record_initial_state( is_async=ctx.background_tasks is not None, worker_pid=worker.get_pid(), worker_started_at=worker.get_started_at(), + dehydrated=bool( + ((query.body or {}).get("parameters") or {}).get("dehydrated") + ), **{ k: v for k, v in query_metadata._asdict().items() diff --git a/src/retriever/utils/mongo.py b/src/retriever/utils/mongo.py index 0c6f67e0..1df0e780 100644 --- a/src/retriever/utils/mongo.py +++ b/src/retriever/utils/mongo.py @@ -2,15 +2,20 @@ import asyncio import base64 +import contextlib import json from collections.abc import AsyncGenerator, Awaitable, Callable -from datetime import datetime +from datetime import datetime, timedelta from typing import Any, ClassVar, Literal, NotRequired, TypedDict, cast, override from bson import ObjectId from bson.codec_options import DEFAULT_CODEC_OPTIONS from loguru import logger as log -from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection +from motor.motor_asyncio import ( + AsyncIOMotorClient, + AsyncIOMotorCollection, + AsyncIOMotorGridFSBucket, +) from pymongo.operations import InsertOne, UpdateOne from pymongo.server_api import ServerApi @@ -22,6 +27,22 @@ CODEC_OPTIONS = DEFAULT_CODEC_OPTIONS.with_options(tz_aware=True) +JOB_DOCS_FS_BUCKET = "job_docs_fs" +"""GridFS bucket name for job-doc blobs too large to store inline.""" + +GRIDFS_INLINE_LIMIT = 8 * 1024 * 1024 +"""Spill a job-doc blob to GridFS past this size. Kept well under Mongo's 16MB +per-document cap so a doc stays valid, and so a flushed batch of inline docs +stays under the ~48MB write-command cap.""" + +GRIDFS_REAP_GRACE = timedelta(minutes=5) +"""Only reap GridFS files older than this, so the sweep can't race a freshly +uploaded blob whose parent `job_docs` document hasn't been written yet.""" + +GRIDFS_REAP_BATCH = 1000 +"""Candidate files resolved per reap batch - bounds the parent-lookup `$in` and +the surviving-id set so neither scales with the total number of stored blobs.""" + _PERCENTILE_MIN_MAJOR = 7 """Mongo major version where `$percentile` became available.""" @@ -53,13 +74,17 @@ class QueryState(TypedDict): worker_pid: int | None worker_started_at: datetime | None event_time: datetime + # True when the query ran dehydrated; recorded at enqueue so the marker + # survives even if the job is abandoned before a response write. + dehydrated: NotRequired[bool] class ResponseState(TypedDict): """Required info about a response for the job state.""" job_id: str - response: bytes + # Omitted for dehydrated queries, whose responses are delivered but never stored. + response: NotRequired[bytes] knodes: int kedges: int aux_graphs: int @@ -67,6 +92,8 @@ class ResponseState(TypedDict): status: str description: NotRequired[str | None] event_time: datetime + # True when the query ran in dehydrated mode; the body is not persisted. + dehydrated: NotRequired[bool] class JobDoc(TypedDict): @@ -75,6 +102,12 @@ class JobDoc(TypedDict): job_id: str doc: NotRequired[bytes] + # Set when the blob was too large to store inline and was spilled to the + # `job_docs_fs` GridFS bucket instead; `doc_ref` is the GridFS file id and + # `get_job_doc` hydrates `doc` from it on read. Mutually exclusive with `doc`. + doc_ref: NotRequired[ObjectId] + doc_size: NotRequired[int] + # Lifetime tracking created: NotRequired[datetime] completed: NotRequired[datetime] @@ -107,6 +140,8 @@ class JobStatus(TypedDict): kedges: NotRequired[int] aux_graphs: NotRequired[int] results: NotRequired[int] + # True when the query ran dehydrated; the response body was deliberately not stored. + dehydrated: NotRequired[bool] # Lifetime tracking created: NotRequired[datetime] @@ -603,18 +638,122 @@ async def ping(self) -> None: """Probe MongoDB. Raises on failure.""" _ = await self.client.admin.command("ping") + def _persist_db(self) -> Any: + """The persistence database. + + Single accessor so the job collections and their GridFS bucket always + resolve to the same DB (and tests can redirect both at once). + """ + return self.client.retriever_persist + def get_job_collection( self, ) -> tuple[ AsyncIOMotorCollection[dict[str, Any]], AsyncIOMotorCollection[dict[str, Any]] ]: """Get job_state collection with standard options.""" - db = self.client.retriever_persist + db = self._persist_db() return ( db.get_collection("job_status", codec_options=CODEC_OPTIONS), db.get_collection("job_docs", codec_options=CODEC_OPTIONS), ) + def _doc_blob_bucket(self) -> AsyncIOMotorGridFSBucket: + """GridFS bucket for job-doc blobs spilled out of `job_docs`. + + Constructed per call - the bucket object is a cheap wrapper over its + backing collections and binds to whatever loop the client is on, so + this avoids caching one against a stale loop after a recovery. + """ + return AsyncIOMotorGridFSBucket( + self._persist_db(), bucket_name=JOB_DOCS_FS_BUCKET + ) + + def _doc_blob_files(self) -> AsyncIOMotorCollection[dict[str, Any]]: + """The GridFS `*.files` collection, for metadata lookups and the reaper.""" + return self._persist_db().get_collection( + f"{JOB_DOCS_FS_BUCKET}.files", codec_options=CODEC_OPTIONS + ) + + async def delete_doc_blobs(self, job_id: str) -> None: + """Delete every GridFS file (and its chunks) tagged with `job_id`.""" + bucket = self._doc_blob_bucket() + async for file in self._doc_blob_files().find( + {"metadata.job_id": job_id}, {"_id": 1} + ): + with contextlib.suppress(Exception): + await bucket.delete(file["_id"]) + + async def offload_doc_blob(self, job_id: str, blob: bytes) -> ObjectId: + """Store an oversized job-doc blob in GridFS and return its file id. + + Drops any prior file for `job_id` first so a query->response overwrite + doesn't leak the earlier upload. + """ + await self.delete_doc_blobs(job_id) + return await self._doc_blob_bucket().upload_from_stream( + job_id, blob, metadata={"job_id": job_id} + ) + + async def download_doc_blob(self, file_id: ObjectId) -> bytes: + """Read a job-doc blob back from GridFS.""" + stream = await self._doc_blob_bucket().open_download_stream(file_id) + return await stream.read() + + async def reap_orphaned_doc_blobs(self) -> int: + """Delete GridFS blobs whose parent `job_docs` document is gone. + + `job_docs` is reaped by Mongo TTL indexes, which don't touch GridFS (and + a TTL index on the files collection would orphan its chunks). This sweep + closes that gap: for files older than `GRIDFS_REAP_GRACE` whose `job_id` + no longer has a `job_docs` document, `bucket.delete` drops the file and + its chunks together. Returns the number of files deleted. + + Candidates are streamed and resolved in fixed-size batches so neither the + parent-lookup `$in` nor the surviving-id set grows with the (potentially + large) number of stored blobs. Only already-streamed files are deleted, + so deleting from the collection being iterated can't skip a candidate. + """ + cutoff = datetime.now().astimezone() - GRIDFS_REAP_GRACE + deleted = 0 + batch: dict[ObjectId, str | None] = {} + async for file in self._doc_blob_files().find( + {"uploadDate": {"$lt": cutoff}}, {"_id": 1, "metadata.job_id": 1} + ): + metadata: dict[str, Any] = file.get("metadata") or {} + batch[file["_id"]] = metadata.get("job_id") + if len(batch) >= GRIDFS_REAP_BATCH: + deleted += await self._reap_doc_blob_batch(batch) + batch = {} + if batch: + deleted += await self._reap_doc_blob_batch(batch) + return deleted + + async def _reap_doc_blob_batch(self, batch: dict[ObjectId, str | None]) -> int: + """Delete the orphans within one batch of candidate GridFS files. + + `batch` maps file id -> the `job_id` it was tagged with. A bounded `$in` + finds which of those jobs still exist; the rest are deleted. + """ + _, docs_collection = self.get_job_collection() + job_ids = {jid for jid in batch.values() if jid is not None} + alive: set[str] = set() + if job_ids: + async for doc in docs_collection.find( + {"job_id": {"$in": sorted(job_ids)}}, {"job_id": 1} + ): + alive.add(doc["job_id"]) + + bucket = self._doc_blob_bucket() + deleted = 0 + for file_id, job_id in batch.items(): + if job_id in alive: + continue + with contextlib.suppress(Exception): + await bucket.delete(file_id) + deleted += 1 + return deleted + def get_log_collection(self) -> AsyncIOMotorCollection[dict[str, Any]]: """Get log_dump collection with standard options.""" db = self.client.retriever_persist @@ -663,6 +802,12 @@ async def _setup_after_connection(self) -> None: ) await collection.create_index("created", background=True) + # Backs the `metadata.job_id` lookups in `delete_doc_blobs` and the + # reaper; the bucket creates its own file/chunk indexes on first upload. + await self._doc_blob_files().create_index( + "metadata.job_id", background=True + ) + log_collection = self.get_log_collection() await log_collection.create_index( {"time": 1}, background=True, expireAfterSeconds=CONFIG.log.mongo_ttl @@ -726,8 +871,20 @@ async def batch_job_state( await self.batch_write(job_status_ops, job_status) await self.batch_write(job_doc_ops, job_docs) - def job_state(self, job: QueryState | ResponseState) -> tuple[UpdateOne, UpdateOne]: - """Create an operation for upserting a job state doc.""" + def job_state( + self, + job: QueryState | ResponseState, + *, + doc_ref: ObjectId | None = None, + store_doc: bool = True, + ) -> tuple[UpdateOne, UpdateOne]: + """Create an operation for upserting a job state doc. + + Pass `doc_ref` when the caller has already spilled the doc blob to + GridFS; the op then stores that reference in place of the inline bytes. + Pass `store_doc=False` when the blob is too large to persist at all; + the op then keeps only metadata and clears any previously stored body. + """ update_time = job["event_time"] status_data = { "$set": JobStatus( @@ -751,12 +908,27 @@ def job_state(self, job: QueryState | ResponseState) -> tuple[UpdateOne, UpdateO } if doc_type == "response": doc_inner["completed"] = update_time - if doc is not None: + # `doc` and `doc_ref` are mutually exclusive; clear the other side so a + # later write (e.g. inline query then GridFS response) leaves no stale field. + unset: dict[str, str] = {} + if not store_doc: + # Blob too large to persist: keep metadata only and drop any prior body. + unset = {"doc": "", "doc_ref": "", "doc_size": ""} + elif doc_ref is not None: + doc_inner["doc_ref"] = doc_ref + if doc is not None: + doc_inner["doc_size"] = len(doc) + unset["doc"] = "" + elif doc is not None: doc_inner["doc"] = doc - doc_data = { + unset["doc_ref"] = "" + unset["doc_size"] = "" + doc_data: dict[str, Any] = { "$set": doc_inner, "$setOnInsert": {"created": update_time}, } + if unset: + doc_data["$unset"] = unset return ( UpdateOne( {"job_id": job["job_id"]}, @@ -799,6 +971,11 @@ async def get_job_doc(self, job_id: str) -> JobDoc | None: ) del job["_id"] + # Blob was spilled to GridFS: hydrate `doc` so callers (`unpack_doc`) + # see the same shape as an inline doc. + doc_ref = job.get("doc_ref") + if doc_ref is not None and job.get("doc") is None: + job["doc"] = await self.download_doc_blob(doc_ref) return JobDoc(**job) async def get_job_status( @@ -1603,12 +1780,43 @@ async def wrapup(self) -> None: ) async def job_state(self, batch: list[QueryState | ResponseState]) -> None: - """Send a batch of job states to MongoDB.""" + """Send a batch of job states to MongoDB. + + Blobs larger than `GRIDFS_INLINE_LIMIT` are spilled to GridFS and stored + by reference, keeping every `job_docs` document under Mongo's cap. Blobs + at or above `CONFIG.mongo.max_stored_doc_bytes`, and dehydrated-query + responses, are not persisted at all - the response is still delivered + upstream, only the stored copy is dropped. + """ if len(batch) == 0: return - status_ops, doc_ops = map( - list, zip(*(self.client.job_state(state) for state in batch), strict=True) - ) + ops: list[tuple[UpdateOne, UpdateOne]] = [] + for state in batch: + doc_type = "query" if "query" in state else "response" + blob = cast("bytes | None", state.get(doc_type)) + size = len(blob) if blob is not None else 0 + dehydrated = doc_type == "response" and bool( + cast("ResponseState", state).get("dehydrated") + ) + doc_ref: ObjectId | None = None + store_doc = True + if dehydrated: + # Dehydrated queries deliver their response but never persist it. + store_doc = False + elif size >= CONFIG.mongo.max_stored_doc_bytes: + store_doc = False + # Drop any earlier body for this job so nothing stale is served back. + await self.client.delete_doc_blobs(state["job_id"]) + log.warning( + f"Job {state['job_id']} {doc_type} is {size} bytes (>= {CONFIG.mongo.max_stored_doc_bytes} cap); delivering without storing it.", + no_mongo_log=True, + ) + elif size > GRIDFS_INLINE_LIMIT: + doc_ref = await self.client.offload_doc_blob(state["job_id"], blob) # pyright: ignore[reportArgumentType] size>0 implies blob is not None + ops.append( + self.client.job_state(state, doc_ref=doc_ref, store_doc=store_doc) + ) + status_ops, doc_ops = map(list, zip(*ops, strict=True)) await self.client.batch_job_state(status_ops, doc_ops) async def log_dump(self, batch: list[dict[str, Any]]) -> None: diff --git a/src/retriever/utils/orphan_detection.py b/src/retriever/utils/orphan_detection.py index 7b72ab5e..0835b87c 100644 --- a/src/retriever/utils/orphan_detection.py +++ b/src/retriever/utils/orphan_detection.py @@ -9,6 +9,8 @@ As a fallback, any non-terminal job older than `CONFIG.job.orphan_max_age` is treated as dead regardless of worker info. We can assume no job legitimately runs that long. +The same loop also reaps GridFS job-doc blobs whose parent `job_docs` +document has been TTL-expired (which Mongo's TTL monitor can't do for us). """ from __future__ import annotations @@ -96,8 +98,24 @@ async def _mark_orphaned_jobs() -> None: logger.exception("Orphan sweep failed to write Failed status.") +async def _reap_orphaned_doc_blobs() -> None: + """One sweep: delete GridFS job-doc blobs whose parent job is gone.""" + if not MongoClient().up: + return + try: + count = await MongoClient().reap_orphaned_doc_blobs() + except Exception: + logger.exception("GridFS blob reap failed.") + return + if count: + logger.info( + f"GridFS reap: deleted {count} orphaned job-doc blob(s).", + no_mongo_log=True, + ) + + async def periodically_mark_orphans() -> None: - """Periodic driver for orphan detection. + """Periodic driver for orphan detection and GridFS blob reaping. Runs immediately on startup so jobs left behind by a previous run get cleaned up without waiting a full interval. There's no startup race: @@ -110,6 +128,7 @@ async def periodically_mark_orphans() -> None: return while True: await _mark_orphaned_jobs() + await _reap_orphaned_doc_blobs() await asyncio.sleep(interval) except asyncio.CancelledError: return diff --git a/tests/test_utils_mongo_gridfs_live.py b/tests/test_utils_mongo_gridfs_live.py new file mode 100644 index 00000000..cd89d5f1 --- /dev/null +++ b/tests/test_utils_mongo_gridfs_live.py @@ -0,0 +1,302 @@ +"""Live tests for GridFS spillover of oversized job-doc blobs. + +Run with: `pytest -m live tests/test_utils_mongo_gridfs_live.py`. Requires a +reachable Mongo (see the docker compose container). Uses the throwaway test DB +from the shared `test_mongo` fixture, which also redirects the GridFS bucket, +so everything is dropped on teardown. +""" + +from __future__ import annotations + +import uuid +from datetime import datetime, timedelta + +import pytest +from utils.mongo_fixtures import ( + test_mongo, # noqa: F401 pyright:ignore[reportImplicitRelativeImport] # fixture import +) + +from retriever.config.general import CONFIG +from retriever.utils.mongo import ( + GRIDFS_INLINE_LIMIT, + JOB_DOCS_FS_BUCKET, + MongoClient, + MongoQueue, + QueryState, + ResponseState, +) + +pytestmark = pytest.mark.live + + +def _response_state(job_id: str, payload: bytes) -> ResponseState: + """A minimal ResponseState carrying `payload` as the stored doc blob.""" + return ResponseState( + job_id=job_id, + response=payload, + knodes=1, + kedges=1, + aux_graphs=0, + results=1, + status="Success", + ) + + +async def _write(state: ResponseState) -> None: + """Persist one state through the queue path that decides inline vs GridFS.""" + await MongoQueue().job_state([state]) + + +@pytest.mark.asyncio +async def test_large_blob_round_trips_via_gridfs(test_mongo: MongoClient) -> None: # noqa: F811 + """A blob over the inline limit is spilled to GridFS and hydrated back intact.""" + job_id = uuid.uuid4().hex + payload = b"\xa5" * (GRIDFS_INLINE_LIMIT + 4096) + await _write(_response_state(job_id, payload)) + + _, docs = test_mongo.get_job_collection() + raw = await docs.find_one({"job_id": job_id}) + assert raw is not None + assert raw.get("doc") is None # not stored inline + assert raw.get("doc_ref") is not None # spilled to GridFS + assert raw.get("doc_size") == len(payload) + + files = test_mongo._doc_blob_files() + assert await files.count_documents({"metadata.job_id": job_id}) == 1 + + job = await test_mongo.get_job_doc(job_id) + assert job is not None + assert job["doc"] == payload # hydrated from GridFS, byte-for-byte + + +@pytest.mark.asyncio +async def test_small_blob_stored_inline(test_mongo: MongoClient) -> None: # noqa: F811 + """A blob under the inline limit stays in the document; no GridFS file.""" + job_id = uuid.uuid4().hex + payload = b"small-response-blob" + await _write(_response_state(job_id, payload)) + + _, docs = test_mongo.get_job_collection() + raw = await docs.find_one({"job_id": job_id}) + assert raw is not None + assert raw.get("doc") == payload + assert raw.get("doc_ref") is None + + files = test_mongo._doc_blob_files() + assert await files.count_documents({"metadata.job_id": job_id}) == 0 + + job = await test_mongo.get_job_doc(job_id) + assert job is not None + assert job["doc"] == payload + + +@pytest.mark.asyncio +async def test_overwrite_leaves_single_gridfs_file(test_mongo: MongoClient) -> None: # noqa: F811 + """Re-spilling a job's blob drops the prior file, leaving exactly one.""" + job_id = uuid.uuid4().hex + first = b"\x01" * (GRIDFS_INLINE_LIMIT + 100) + second = b"\x02" * (GRIDFS_INLINE_LIMIT + 200) + await _write(_response_state(job_id, first)) + await _write(_response_state(job_id, second)) + + files = test_mongo._doc_blob_files() + assert await files.count_documents({"metadata.job_id": job_id}) == 1 + + job = await test_mongo.get_job_doc(job_id) + assert job is not None + assert job["doc"] == second + + +@pytest.mark.asyncio +async def test_reaper_deletes_orphaned_blob( + test_mongo: MongoClient, # noqa: F811 + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The reaper drops a blob (and its chunks) whose job_docs doc is gone.""" + # Drop the grace window so a just-uploaded file is immediately eligible. + monkeypatch.setattr("retriever.utils.mongo.GRIDFS_REAP_GRACE", timedelta(0)) + + alive = uuid.uuid4().hex + orphan = uuid.uuid4().hex + await _write(_response_state(alive, b"\x03" * (GRIDFS_INLINE_LIMIT + 1))) + await _write(_response_state(orphan, b"\x04" * (GRIDFS_INLINE_LIMIT + 1))) + + # Simulate the orphan's job_docs document being TTL-expired. + _, docs = test_mongo.get_job_collection() + await docs.delete_one({"job_id": orphan}) + + assert await test_mongo.reap_orphaned_doc_blobs() == 1 + + files = test_mongo._doc_blob_files() + chunks = test_mongo._persist_db().get_collection(f"{JOB_DOCS_FS_BUCKET}.chunks") + assert await files.count_documents({"metadata.job_id": orphan}) == 0 + assert await files.count_documents({"metadata.job_id": alive}) == 1 + + # The orphan's chunks are gone too; only the surviving file's remain. + alive_file = await files.find_one({"metadata.job_id": alive}) + assert alive_file is not None + assert await chunks.count_documents({"files_id": {"$ne": alive_file["_id"]}}) == 0 + assert await chunks.count_documents({"files_id": alive_file["_id"]}) > 0 + + +@pytest.mark.asyncio +async def test_reaper_streams_across_batches( + test_mongo: MongoClient, # noqa: F811 + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reaping is correct when candidates span multiple batches (batch size 1).""" + monkeypatch.setattr("retriever.utils.mongo.GRIDFS_REAP_GRACE", timedelta(0)) + monkeypatch.setattr("retriever.utils.mongo.GRIDFS_REAP_BATCH", 1) + + alive = uuid.uuid4().hex + orphans = [uuid.uuid4().hex for _ in range(2)] + for job_id in (alive, *orphans): + await _write(_response_state(job_id, b"\x05" * (GRIDFS_INLINE_LIMIT + 1))) + + _, docs = test_mongo.get_job_collection() + await docs.delete_many({"job_id": {"$in": orphans}}) + + assert await test_mongo.reap_orphaned_doc_blobs() == 2 + + files = test_mongo._doc_blob_files() + assert await files.count_documents({"metadata.job_id": {"$in": orphans}}) == 0 + assert await files.count_documents({"metadata.job_id": alive}) == 1 + + +@pytest.mark.asyncio +async def test_blob_at_cap_is_delivered_but_not_stored( + test_mongo: MongoClient, # noqa: F811 + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A blob at/above the cap leaves a terminal status but no stored body.""" + cap = 512 * 1024 + monkeypatch.setattr(CONFIG.mongo, "max_stored_doc_bytes", cap) + + job_id = uuid.uuid4().hex + await _write(_response_state(job_id, b"\x06" * (cap + 1024))) + + status, docs = test_mongo.get_job_collection() + # Status is still recorded, so the job reads back as a completed job. + status_doc = await status.find_one({"job_id": job_id}) + assert status_doc is not None + assert status_doc["status"] == "Success" + + # The doc carries no body in any form, and nothing was spilled to GridFS. + raw = await docs.find_one({"job_id": job_id}) + assert raw is not None + assert raw.get("doc") is None + assert raw.get("doc_ref") is None + assert raw.get("doc_size") is None + assert ( + await test_mongo._doc_blob_files().count_documents({"metadata.job_id": job_id}) + == 0 + ) + + job = await test_mongo.get_job_doc(job_id) + assert job is not None + assert job.get("doc") is None + + +@pytest.mark.asyncio +async def test_cap_clears_a_previously_stored_body( + test_mongo: MongoClient, # noqa: F811 + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An over-cap rewrite drops the body stored by an earlier under-cap write.""" + job_id = uuid.uuid4().hex + # First write is small enough to store inline. + await _write(_response_state(job_id, b"first-small-body")) + + _, docs = test_mongo.get_job_collection() + assert (await docs.find_one({"job_id": job_id}) or {}).get("doc") is not None + + # Second write is over the cap, so the prior body must be cleared. + cap = 512 * 1024 + monkeypatch.setattr(CONFIG.mongo, "max_stored_doc_bytes", cap) + await _write(_response_state(job_id, b"\x07" * (cap + 1024))) + + raw = await docs.find_one({"job_id": job_id}) + assert raw is not None + assert raw.get("doc") is None + assert raw.get("doc_ref") is None + + +@pytest.mark.asyncio +async def test_dehydrated_response_stores_state_but_no_body( + test_mongo: MongoClient, # noqa: F811 +) -> None: + """Dehydrated queries persist their state (counts/status) but never a body.""" + job_id = uuid.uuid4().hex + # An earlier inline write (e.g. the initial query state) leaves a body... + await _write(_response_state(job_id, b"earlier-inline-body")) + status, docs = test_mongo.get_job_collection() + assert (await docs.find_one({"job_id": job_id}) or {}).get("doc") is not None + + # ...which the dehydrated response write must clear, keeping state only. + await MongoQueue().job_state( + [ + ResponseState( + job_id=job_id, + knodes=3, + kedges=2, + aux_graphs=0, + results=5, + status="Success", + dehydrated=True, + ) + ] + ) + + status_doc = await status.find_one({"job_id": job_id}) + assert status_doc is not None + assert status_doc["status"] == "Success" + assert status_doc.get("dehydrated") is True + assert status_doc.get("results") == 5 # state is recorded + + raw = await docs.find_one({"job_id": job_id}) + assert raw is not None + assert raw.get("doc") is None + assert raw.get("doc_ref") is None + assert ( + await test_mongo._doc_blob_files().count_documents({"metadata.job_id": job_id}) + == 0 + ) + + job = await test_mongo.get_job_doc(job_id) + assert job is not None + assert job.get("doc") is None + + +@pytest.mark.asyncio +async def test_dehydrated_flag_recorded_at_enqueue(test_mongo: MongoClient) -> None: # noqa: F811 + """The initial query write records `dehydrated` and still stores the query body.""" + job_id = uuid.uuid4().hex + await MongoQueue().job_state( + [ + QueryState( + job_id=job_id, + query=b"compressed-query-bytes", + job_timeout=30.0, + submitter="tester", + data_tier=0, + is_async=True, + qnodes=2, + qedges=1, + qpaths=0, + status="Running", + worker_pid=1234, + worker_started_at=datetime.now().astimezone(), + dehydrated=True, + ) + ] + ) + + status, docs = test_mongo.get_job_collection() + status_doc = await status.find_one({"job_id": job_id}) + assert status_doc is not None + assert status_doc.get("dehydrated") is True # marker survives even if abandoned + + # Dehydrated only skips response bodies; the query itself is still stored. + raw = await docs.find_one({"job_id": job_id}) + assert raw is not None + assert raw.get("doc") == b"compressed-query-bytes" diff --git a/tests/utils/mongo_fixtures.py b/tests/utils/mongo_fixtures.py index 85056df4..92e758c1 100644 --- a/tests/utils/mongo_fixtures.py +++ b/tests/utils/mongo_fixtures.py @@ -44,6 +44,9 @@ def patched_get_job_collection() -> ( ) monkeypatch.setattr(client, "get_job_collection", patched_get_job_collection) + # Redirect the GridFS bucket (job-doc spillover) at the same throwaway DB so + # large-blob tests don't write into the real persistence DB. + monkeypatch.setattr(client, "_persist_db", lambda: test_db) status_collection, _ = patched_get_job_collection() await status_collection.create_index("job_id", unique=True, background=True)