diff --git a/hindsight-api-slim/hindsight_api/engine/maintenance.py b/hindsight-api-slim/hindsight_api/engine/maintenance.py index 772b87ff66..fb78423b11 100644 --- a/hindsight-api-slim/hindsight_api/engine/maintenance.py +++ b/hindsight-api-slim/hindsight_api/engine/maintenance.py @@ -25,6 +25,12 @@ ``banks_needing_consolidation``, in the configured schema — see ``fq_routine``) — one round-trip each — instead of a per-schema query storm, which matters at thousands of tenants. + +The loop runs in *every* API/worker process with no leader election, so a job that +enqueues work must make that enqueue idempotent or the fleet queues one wave per +process. Retention and operation cleanup are deletes; the consolidation reconcile +and the scheduled mental model refresh both dedupe against in-flight operations +inside the inserting transaction (see ``_submit_async_operation``). """ from __future__ import annotations @@ -450,6 +456,7 @@ async def _run_scheduled_mm_refresh(self) -> None: submitted = 0 skipped_unknown = 0 skipped_fresh = 0 + skipped_in_flight = 0 for row in due: schema = row["schema_name"] bank_id = row["bank_id"] @@ -480,18 +487,28 @@ async def _run_scheduled_mm_refresh(self) -> None: if not is_stale: skipped_fresh += 1 continue - await engine.submit_async_refresh_mental_model( - bank_id=bank_id, mental_model_id=mm_id, request_context=context + # skip_if_in_flight makes the enqueue itself idempotent. The discovery + # routine already excludes models with a pending/processing refresh, + # but that exclusion is a *read*: this loop runs in every process, so + # every process saw the same "nothing in flight" snapshot and inserted + # its own operation — one queued wave per process (#3210). The insert + # now carries the check, so a second one is never created. + result = await engine.submit_async_refresh_mental_model( + bank_id=bank_id, mental_model_id=mm_id, request_context=context, skip_if_in_flight=True ) - submitted += 1 + if result.get("deduplicated"): + skipped_in_flight += 1 + else: + submitted += 1 except Exception as e: logger.warning(f"Scheduled mental model refresh failed for {mm_id} in {schema}: {e}") finally: _current_schema.reset(token) - if submitted or skipped_unknown or skipped_fresh: + if submitted or skipped_unknown or skipped_fresh or skipped_in_flight: logger.info( f"Scheduled mental model refresh: scheduled {submitted} model(s)" + (f", {skipped_fresh} up-to-date" if skipped_fresh else "") + + (f", {skipped_in_flight} already in flight" if skipped_in_flight else "") + (f", skipped {skipped_unknown} in unrecognized schema(s)" if skipped_unknown else "") ) diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index db9924f383..5473cfe4aa 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -14853,6 +14853,7 @@ async def _submit_async_operation( *, result_metadata: dict[str, Any] | None = None, dedupe_by_bank: bool = False, + dedupe_in_flight_payload_key: str | None = None, ) -> dict[str, Any]: """Generic helper to submit an async operation. @@ -14863,6 +14864,9 @@ async def _submit_async_operation( task_payload: Additional task payload fields (operation_id and bank_id are added automatically) result_metadata: Optional metadata to store with the operation record dedupe_by_bank: If True, skip creating a new task if one is already pending for this bank+operation_type + dedupe_in_flight_payload_key: If set, skip creating a new task when a pending or processing + operation of this type exists whose task_payload carries the same value for this key + (e.g. 'mental_model_id'). Narrower than dedupe_by_bank, which dedupes per bank. Returns: Dict with operation_id and optionally deduplicated=True if an existing task was found @@ -14888,39 +14892,42 @@ async def _submit_async_operation( async with acquire_with_retry(backend) as conn: async with conn.transaction(): + # Serialize concurrent submits for this bank whenever the INSERT is + # conditional on what is already queued, so the check-and-insert is + # atomic. A bare check-then-INSERT races under READ COMMITTED: two + # /consolidate calls (or a manual trigger racing a retain-driven + # submit / round-limit re-queue) both see no pending row and both + # insert, leaking duplicate pending ops that then pile up as + # retry_blocked and starve the bank (issue #1842). Locking the bank + # row serializes submits for this bank; it releases on commit below. + # + # FOR NO KEY UPDATE, not FOR UPDATE: async_operations has an FK to + # banks, so every async-op insert for this bank (a scoped + # consolidation, a batch-retain op, a webhook delivery, ...) takes a + # FOR KEY SHARE lock on the bank row. FOR UPDATE conflicts with + # FOR KEY SHARE and would block all of those during the submit; + # FOR NO KEY UPDATE still conflicts with itself (so two submits + # serialize) but not with FOR KEY SHARE (so those inserts proceed). + # On Oracle this rewrites to FOR UPDATE, which there does not block + # indexed-FK child inserts. + # + # Unconditional submits skip the lock but still verify the bank + # exists: without the check, callers that race against bank deletion + # or that derive bank IDs before creating the bank reach the INSERT + # below and get an asyncpg.ForeignKeyViolationError, which surfaces + # as an opaque 500 from the API. A clean OperationValidationError(404) + # is the right shape — the FastAPI handler already converts it via its + # existing except clause. + serialize = dedupe_by_bank or dedupe_in_flight_payload_key is not None + bank_exists = await conn.fetchval( + f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1" + + (" FOR NO KEY UPDATE" if serialize else ""), + bank_id, + ) + if bank_exists is None: + raise OperationValidationError(f"Bank '{bank_id}' not found", status_code=404) + if dedupe_by_bank: - # Serialize concurrent submits for this bank so the dedup - # check-and-insert is atomic. A bare check-then-INSERT races - # under READ COMMITTED: two /consolidate calls (or a manual - # trigger racing a retain-driven submit / round-limit re-queue) - # both see no pending row and both insert, leaking duplicate - # pending ops that then pile up as retry_blocked and starve the - # bank (issue #1842). Locking the bank row serializes submits for - # this bank; it releases on commit below. - # - # FOR NO KEY UPDATE, not FOR UPDATE: async_operations has an FK to - # banks, so every async-op insert for this bank (a scoped - # consolidation, a batch-retain op, a webhook delivery, ...) takes a - # FOR KEY SHARE lock on the bank row. FOR UPDATE conflicts with - # FOR KEY SHARE and would block all of those during the submit; - # FOR NO KEY UPDATE still conflicts with itself (so two submits - # serialize) but not with FOR KEY SHARE (so those inserts proceed). - # On Oracle this rewrites to FOR UPDATE, which there does not block - # indexed-FK child inserts. - # - # Use fetchval so we can also verify the bank actually exists. - # Without this check, callers that race against bank deletion - # or that derive bank IDs before creating the bank reach the - # INSERT below and get an asyncpg.ForeignKeyViolationError, which - # surfaces as an opaque 500 from the API. A clean - # OperationValidationError(404) is the right shape — the FastAPI - # handler already converts it via its existing except clause. - bank_exists = await conn.fetchval( - f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1 FOR NO KEY UPDATE", - bank_id, - ) - if bank_exists is None: - raise OperationValidationError(f"Bank '{bank_id}' not found", status_code=404) # Only check 'pending', not 'processing': a processing task uses a # watermark from when it started, so memories added after that need # a fresh run regardless. @@ -14950,22 +14957,7 @@ async def _submit_async_operation( "operation_id": str(row["operation_id"]), "deduplicated": True, } - else: - # Scoped/non-dedupe submits skip the lock + dedup above. - # Still verify the bank exists so an FK violation can't - # escape as a 500. - bank_exists = await conn.fetchval( - f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1", - bank_id, - ) - if bank_exists is None: - raise OperationValidationError(f"Bank '{bank_id}' not found", status_code=404) - - await conn.execute( - f""" - INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status, task_payload) - VALUES ($1, $2, $3, $4, $5, $6::jsonb) - """, + insert_args = ( operation_id, bank_id, operation_type, @@ -14973,6 +14965,69 @@ async def _submit_async_operation( "pending", json.dumps(full_payload, default=_json_default), ) + if dedupe_in_flight_payload_key is None: + await conn.execute( + f""" + INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status, task_payload) + VALUES ($1, $2, $3, $4, $5, $6::jsonb) + """, + *insert_args, + ) + else: + # Sub-bank dedup: the INSERT itself only materialises a row when no + # operation of this type is already queued or running for the same + # payload subject (e.g. one mental model), so the check cannot be + # separated from the write (#3210). + # + # 'processing' counts here, unlike the bank-wide branch above: the + # only caller is the cron-scheduled refresh, whose next tick covers + # anything the in-flight run misses, so a second op would just + # re-check staleness and occupy a claim slot. + # + # PostgreSQL JSON syntax: this path is reached only from the + # maintenance loop, which is PostgreSQL-only. Oracle submits take + # the unconditional branch above. + subject = task_payload.get(dedupe_in_flight_payload_key) + inserted = await conn.fetchval( + f""" + INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status, task_payload) + SELECT $1::uuid, $2, $3, $4::jsonb, $5::text, $6::jsonb + WHERE NOT EXISTS ( + SELECT 1 FROM {fq_table("async_operations")} + WHERE bank_id = $2 AND operation_type = $3 + AND status IN ('pending', 'processing') + AND task_payload->>$7 = $8 + ) + RETURNING operation_id + """, + *insert_args, + dedupe_in_flight_payload_key, + subject, + ) + if inserted is None: + existing = await conn.fetchval( + f""" + SELECT operation_id FROM {fq_table("async_operations")} + WHERE bank_id = $1 AND operation_type = $2 + AND status IN ('pending', 'processing') + AND task_payload->>$3 = $4 + ORDER BY created_at + LIMIT 1 + """, + bank_id, + operation_type, + dedupe_in_flight_payload_key, + subject, + ) + logger.debug( + f"{operation_type} task already in flight for bank_id={bank_id} " + f"{dedupe_in_flight_payload_key}={subject}, skipping duplicate " + f"(existing operation_id={existing})" + ) + return { + "operation_id": str(existing), + "deduplicated": True, + } # For SyncTaskBackend: executes the task immediately. # For BrokerTaskBackend: no-op (submit_task's UPDATE skips rows whose @@ -15500,6 +15555,7 @@ async def submit_async_refresh_mental_model( mental_model_id: str, *, request_context: "RequestContext", + skip_if_in_flight: bool = False, ) -> dict[str, Any]: """Submit an async mental model refresh operation. @@ -15509,6 +15565,12 @@ async def submit_async_refresh_mental_model( bank_id: Bank identifier mental_model_id: Mental model UUID to refresh request_context: Request context for authentication + skip_if_in_flight: If True, return the existing operation (with + ``deduplicated=True``) instead of queueing a second refresh when one is + already pending or processing for this model. Used by the scheduled + (cron) refresh, which runs in every process of the fleet and would + otherwise queue one wave per process (#3210). Explicit user-triggered + refreshes leave it False so an on-demand refresh is never swallowed. Returns: Dict with operation_id @@ -15556,6 +15618,7 @@ async def submit_async_refresh_mental_model( task_payload=task_payload, result_metadata={"mental_model_id": mental_model_id, "name": mental_model["name"]}, dedupe_by_bank=False, + dedupe_in_flight_payload_key="mental_model_id" if skip_if_in_flight else None, ) def _raise_if_mental_model_refresh_unavailable(self) -> None: diff --git a/hindsight-api-slim/tests/test_mental_model_delta.py b/hindsight-api-slim/tests/test_mental_model_delta.py index 87a694faa5..36f423a8d4 100644 --- a/hindsight-api-slim/tests/test_mental_model_delta.py +++ b/hindsight-api-slim/tests/test_mental_model_delta.py @@ -383,7 +383,11 @@ async def fail_embedding_generation(*args, **kwargs): submitted: list[str] = [] async def record_submit( - *, bank_id: str, mental_model_id: str, request_context: RequestContext + *, + bank_id: str, + mental_model_id: str, + request_context: RequestContext, + skip_if_in_flight: bool = False, ) -> dict[str, str]: submitted.append(mental_model_id) return {"operation_id": str(uuid.uuid4())} diff --git a/hindsight-api-slim/tests/test_mental_model_scheduled_refresh.py b/hindsight-api-slim/tests/test_mental_model_scheduled_refresh.py index f273d714d8..ba93fee9a3 100644 --- a/hindsight-api-slim/tests/test_mental_model_scheduled_refresh.py +++ b/hindsight-api-slim/tests/test_mental_model_scheduled_refresh.py @@ -7,6 +7,7 @@ monkeypatched. """ +import asyncio import json import uuid @@ -75,7 +76,7 @@ async def _insert_fact(conn, bank_id: str, tags: list[str] | None = None) -> Non def _patch_submit(memory: MemoryEngine, monkeypatch) -> list[str]: submitted: list[str] = [] - async def _record(*, bank_id, mental_model_id, request_context): + async def _record(*, bank_id, mental_model_id, request_context, skip_if_in_flight=False): submitted.append(mental_model_id) return {"operation_id": str(uuid.uuid4())} @@ -83,6 +84,27 @@ async def _record(*, bank_id, mental_model_id, request_context): return submitted +def _stall_worker(memory: MemoryEngine, monkeypatch) -> None: + """Queue operations without executing them. + + Reproduces the condition the duplicate waves were observed under (#3210): the + ops stay pending because completions stalled fleet-wide. + """ + + async def _never_runs(task_dict): + return None + + monkeypatch.setattr(memory._task_backend, "submit_task", _never_runs) + + +async def _count_refresh_ops(memory: MemoryEngine, bank_id: str) -> int: + async with memory._pool.acquire() as conn: + return await conn.fetchval( + "SELECT count(*) FROM async_operations WHERE bank_id = $1 AND operation_type = 'refresh_mental_model'", + bank_id, + ) + + @pytest.mark.asyncio async def test_refresh_cron_round_trips_through_create_and_get(memory: MemoryEngine, request_context): """refresh_cron set on a mental model's trigger persists and reads back.""" @@ -158,6 +180,62 @@ async def test_due_but_not_stale_model_is_skipped(memory: MemoryEngine, request_ assert mm_id not in submitted +@pytest.mark.asyncio +async def test_concurrent_scheduled_submits_queue_one_refresh(memory: MemoryEngine, request_context, monkeypatch): + """Two schedulers enqueueing the same due model at once queue exactly one op. + + Regression for #3210: the maintenance loop runs in every process, and the + in-flight guard in ``mental_models_with_cron()`` is a *read* — every process saw + the same "nothing in flight" snapshot and inserted its own operation, so a few + hundred due models became thousands of queued refreshes. The enqueue now does the + in-flight check under the bank row lock, inside the inserting transaction. + """ + bank = await _make_bank(memory, request_context) + async with memory._pool.acquire() as conn: + mm_id = await _insert_mm(conn, bank, refresh_cron="*/5 * * * *", last_refreshed_offset="1 day") + _stall_worker(memory, monkeypatch) + + first, second = await asyncio.gather( + memory.submit_async_refresh_mental_model( + bank_id=bank, mental_model_id=mm_id, request_context=request_context, skip_if_in_flight=True + ), + memory.submit_async_refresh_mental_model( + bank_id=bank, mental_model_id=mm_id, request_context=request_context, skip_if_in_flight=True + ), + ) + + assert await _count_refresh_ops(memory, bank) == 1 + # The loser reports the winner's operation rather than a fresh one, so the + # maintenance loop can count it as "already in flight". + deduped = second if second.get("deduplicated") else first + kept = first if second.get("deduplicated") else second + assert deduped["deduplicated"] is True + assert deduped["operation_id"] == kept["operation_id"] + + +@pytest.mark.asyncio +async def test_user_triggered_refresh_is_not_deduplicated(memory: MemoryEngine, request_context, monkeypatch): + """An explicit refresh still queues while a scheduled one is in flight. + + The dedup is opt-in for the cron scheduler only: a user asking for a refresh has + new intent (e.g. an edited source query) and must not be silently swallowed. + """ + bank = await _make_bank(memory, request_context) + async with memory._pool.acquire() as conn: + mm_id = await _insert_mm(conn, bank, refresh_cron="*/5 * * * *", last_refreshed_offset="1 day") + _stall_worker(memory, monkeypatch) + + scheduled = await memory.submit_async_refresh_mental_model( + bank_id=bank, mental_model_id=mm_id, request_context=request_context, skip_if_in_flight=True + ) + manual = await memory.submit_async_refresh_mental_model( + bank_id=bank, mental_model_id=mm_id, request_context=request_context + ) + + assert manual["operation_id"] != scheduled["operation_id"] + assert await _count_refresh_ops(memory, bank) == 2 + + @pytest.mark.asyncio async def test_not_due_model_is_skipped_even_when_stale(memory: MemoryEngine, request_context, monkeypatch): """A model whose cron has not elapsed since the last refresh is not refreshed,