From 111c7c95b51ae4128afae9c9103b2e8066a9a34d Mon Sep 17 00:00:00 2001 From: shehjad-dev Date: Wed, 30 Sep 2026 06:06:13 +0600 Subject: [PATCH] fix(audit): queue async terminal hit correction --- docs/context/architecture/catalog.md | 3 +- docs/context/architecture/data-model.md | 2 +- src/treg/application/asynctasks.py | 2 +- src/treg/audit.py | 46 +++++++++++-------- tests/test_db_pool_isolation.py | 1 - tests/test_wiza.py | 60 +++++++++++++++++++++++++ 6 files changed, 91 insertions(+), 23 deletions(-) diff --git a/docs/context/architecture/catalog.md b/docs/context/architecture/catalog.md index a2d5b117..3d2f7330 100644 --- a/docs/context/architecture/catalog.md +++ b/docs/context/architecture/catalog.md @@ -1927,7 +1927,8 @@ to choose (`docs/CAPABILITY-ROUTING-PLAN.md`). Everything else in the catalog st failures may continue under the normal bounded fallback rules. - **Hit rate** — `CallRecord.hit` (nullable, alembic `0009`, last column) is the adapter's verdict for a completed result. An accepted async submission leaves it NULL; the terminal finalizer - classifies the finished body and updates the originating row, including a routed child. A + classifies the finished body and queues an audit update for the originating row, including a + routed child. The poll response does not wait for that best-effort write. A confirmed terminal failure counts as a miss; pending and timed-out jobs remain undecided. Its `AsyncTaskRecord.hit` keeps the verdict if polling beats the background audit insert. Async endpoints read their `CallRecord` observations live: the daily fold may consume a diff --git a/docs/context/architecture/data-model.md b/docs/context/architecture/data-model.md index 53d0fa0d..a57b3437 100644 --- a/docs/context/architecture/data-model.md +++ b/docs/context/architecture/data-model.md @@ -90,7 +90,7 @@ authorize a shared-provider result without them. Revision `0052` adds nullable `hit`: the terminal contact verdict for the original submission. The terminal finalizer commits it with settlement. The audit writer locks the task row before inserting a submission `CallRecord`, so a poll that finishes first still gives that row its final -verdict; when the audit row wins the race, the finalizer updates it afterward. Both use the +verdict; when the audit row wins the race, the finalizer queues a background correction. Both use the original `call_ref`, including routed children. A confirmed terminal failure stores `false` for endpoints with verified result rules, since that attempt produced no hit. A pending or timed-out submission remains undecided. diff --git a/src/treg/application/asynctasks.py b/src/treg/application/asynctasks.py index aa89132d..8754c577 100644 --- a/src/treg/application/asynctasks.py +++ b/src/treg/application/asynctasks.py @@ -403,7 +403,7 @@ async def _finish_terminal(snapshot: AsyncTaskRecord, outcome: str, document: ob expected = asynctasks.SETTLED if outcome in ("success", "billed_failure") else asynctasks.RELEASED if result == expected: if terminal_hit is not None: - await audit.update_async_call_hit( + audit.record_async_call_hit( snapshot.call_id, snapshot.endpoint_id, snapshot.org_id, terminal_hit) # Only the winning finalizer records evidence; a late poll cannot replace the result # whose usage was charged. Archive failure cannot undo the committed money transaction. diff --git a/src/treg/audit.py b/src/treg/audit.py index ce8d07aa..052d32db 100644 --- a/src/treg/audit.py +++ b/src/treg/audit.py @@ -1,6 +1,7 @@ """Audit writes — deferred and fire-and-forget (rule #2: never block the proxied response). -`record_call` queues a row and returns immediately; the response streams without waiting. One +`record_call` and terminal hit corrections queue their writes and return immediately; the +response streams without waiting. One writer task per process drains the queue in batches on one connection (a strong reference to it is held until it finishes, otherwise the event loop may GC a bare create_task). Failures are swallowed: an audit hiccup must never break a real call. `drain()` flushes pending writes on @@ -9,7 +10,7 @@ shutdown / in tests. Back-pressure (why this matters): the writer's connection comes from the BACKGROUND pool (db.py), so a burst here can starve other background work but never real calls. Rows queue in-process, not as pooled connections: one writer per process takes them off the queue `_BATCH` at a time and lands -each batch in one INSERT round trip, which is what keeps `drain()` deterministic on sqlite, where +inserts in batches, which is what keeps `drain()` deterministic on sqlite, where all three makers share one engine. Under an extreme burst we DROP audit rows past `_MAX_PENDING` rather than grow without bound — audit is best-effort; never OOM or wedge the server for it. """ @@ -33,9 +34,14 @@ _pending: set[asyncio.Task] = set() # lands the same rows in fewer round trips and holds one connection. _MAX_CONCURRENT_WRITES = 1 _MAX_PENDING = 5000 # shed load past this: drop the audit row rather than grow unbounded -_BATCH = 200 # rows per INSERT round trip; a failed batch retries row by row +_BATCH = 200 # queued operations per batch; a failed batch retries one by one _queue: deque[tuple[type, dict]] = deque() + +class _AsyncHitUpdate: + """Queued audit mutation, ordered with inserts on the same writer.""" + + _sem: asyncio.Semaphore | None = None _sem_loop = None @@ -71,24 +77,14 @@ def record_call( )) -async def update_async_call_hit(call_id: str, endpoint_id: str, org_id: int, hit: bool) -> None: - """Copy a terminal verdict to the original audit row, if it has been written. +def record_async_call_hit(call_id: str, endpoint_id: str, org_id: int, hit: bool) -> None: + """Queue the terminal audit correction without delaying the provider's response. - The task row is the durable source when the audit insert has not happened yet. Its - row lock also orders this update against a concurrent audit writer. Share the audit - writer's background-pool slot rather than creating another concurrent consumer. + If the insert has not landed, its task-row read supplies the durable verdict. If it + has landed, this update corrects it. Both operations use the one audit writer. """ - try: - async with _get_sem(): - async with background_session_maker() as session: - await session.execute(update(CallRecord).where( - CallRecord.call_ref == call_id, - CallRecord.endpoint_id == endpoint_id, - CallRecord.org_id == org_id, - ).values(hit=hit)) - await session.commit() - except Exception: # noqa: BLE001 — audit cannot undo a committed settlement - logging.getLogger("treg.audit").error("async hit update failed for %s", call_id, exc_info=True) + _enqueue(_AsyncHitUpdate, dict(call_id=call_id, endpoint_id=endpoint_id, + org_id=org_id, hit=hit)) def _known_fields(model, telemetry: dict | None) -> dict: @@ -213,6 +209,18 @@ async def _write_batch(rows: list[tuple[type, dict]]) -> bool: .with_for_update())).scalars()} records = [] for model, fields in rows: + if model is _AsyncHitUpdate: + # Earlier inserts in this batch must be visible to the UPDATE. + if records: + session.add_all(records) + records = [] + await session.flush() + await session.execute(update(CallRecord).where( + CallRecord.call_ref == fields["call_id"], + CallRecord.endpoint_id == fields["endpoint_id"], + CallRecord.org_id == fields["org_id"], + ).values(hit=fields["hit"])) + continue values = {k: v for k, v in fields.items() if k != "_async_submission"} if fields.get("_async_submission") and (task := tasks.get(fields["call_ref"])) is not None: values["hit"] = task.hit diff --git a/tests/test_db_pool_isolation.py b/tests/test_db_pool_isolation.py index 09c0d8e3..50aeb990 100644 --- a/tests/test_db_pool_isolation.py +++ b/tests/test_db_pool_isolation.py @@ -171,7 +171,6 @@ BACKGROUND_SITES = { "bootstrap.py:_lifespan.lifespan": "adsconv.worker", "bootstrap.py:create_app": "catalog observation refresh", "audit.py:_write_batch": "audit._flush", - "audit.py:update_async_call_hit": "audit._flush", "archive.py:_store_locked": "archive._store/_touch", "archive.py:_touch_write": "archive._store/_touch", "archive.py:_ignored_matches": "archive._store/_touch", diff --git a/tests/test_wiza.py b/tests/test_wiza.py index d498cdaa..24cca6d4 100644 --- a/tests/test_wiza.py +++ b/tests/test_wiza.py @@ -2,7 +2,9 @@ from __future__ import annotations +import asyncio import json +from contextlib import asynccontextmanager import httpx import pytest @@ -247,6 +249,64 @@ async def test_wiza_terminal_hit_precedes_audit_insert(clients, monkeypatch, wiz assert original.hit is True +async def test_wiza_terminal_poll_does_not_wait_for_audit_writer( + clients, monkeypatch, wiza_platform_on, +): + await audit.drain() + writer_entered = asyncio.Event() + release_writer = asyncio.Event() + real_session_maker = audit.background_session_maker + + @asynccontextmanager + async def gated_session_maker(): + writer_entered.set() + await release_writer.wait() + async with real_session_maker() as session: + yield session + + monkeypatch.setattr(audit, "background_session_maker", gated_session_maker) + + async def relay(request, upstream_url, tool, secrets, client, **kwargs): + doc = ({"data": {"id": 5432, "status": "queued"}} if request.method == "POST" else + {"data": {"id": 5432, "status": "finished", "email": "person@sample.example", + "email_status": "risky", "credits": {"api_credits": {"total": 0}}}}) + payload = json.dumps(doc).encode() + + async def stream(): + yield payload + + async def close(): + return None + + return UpstreamResponse(200, ((b"content-type", b"application/json"),), stream(), close) + + monkeypatch.setattr(call_service, "relay", relay) + call_ref = None + try: + submitted = await asyncio.wait_for(clients.post("/call/wiza.people.email.find", json={ + "individual_reveal": {"full_name": "Person Example", "domain": "sample.example"}, + "enrichment_level": "partial", + "email_options": {"accept_work": True, "accept_personal": False, + "accept_generic": False}, + }), timeout=5) + assert submitted.status_code == 200 + call_ref = submitted.headers["x-treg-call-id"] + await asyncio.wait_for(writer_entered.wait(), timeout=5) + polled = await asyncio.wait_for(clients.get( + "/call/wiza.people.reveal.get", params={"id": 5432}), timeout=5) + assert polled.status_code == 200 + async with session_maker() as db: + assert (await db.get(AsyncTaskRecord, call_ref)).hit is True + finally: + release_writer.set() + await audit.drain() + async with session_maker() as db: + original = (await db.execute(select(CallRecord).where( + CallRecord.call_ref == call_ref, + CallRecord.endpoint_id == "wiza.people.email.find"))).scalar_one() + assert original.hit is True + + @pytest.mark.parametrize("capability,terminal,expected_hit", [ ("email", {"email": "person@sample.example", "email_status": "risky"}, True), ("phone", {"phone_status": "unfound"}, False),