From 6dba0fa8f12b8ca71c2ebe628f52df749a6b0b90 Mon Sep 17 00:00:00 2001 From: ogt Date: Wed, 15 Jul 2026 01:41:09 +0800 Subject: [PATCH] fix(api): coalesce consumer readback DB queries --- ..._controlled_writeback_consumer_readback.py | 115 +++++++++++++++++- ...trolled_writeback_consumer_readback_api.py | 101 +++++++++++++++ 2 files changed, 215 insertions(+), 1 deletion(-) diff --git a/apps/api/src/services/ai_agent_log_controlled_writeback_consumer_readback.py b/apps/api/src/services/ai_agent_log_controlled_writeback_consumer_readback.py index 040dd5ff2..501958458 100644 --- a/apps/api/src/services/ai_agent_log_controlled_writeback_consumer_readback.py +++ b/apps/api/src/services/ai_agent_log_controlled_writeback_consumer_readback.py @@ -11,8 +11,11 @@ prove this route has rolled out. from __future__ import annotations import asyncio +import copy import json +import time from collections.abc import Mapping +from datetime import UTC, datetime from typing import Any from sqlalchemy import text @@ -51,6 +54,13 @@ _CONSUMER_DIRECT_CONNECT_TIMEOUT_SECONDS = 1.5 _CONSUMER_DIRECT_QUERY_TIMEOUT_SECONDS = 2.0 _CONSUMER_TELEGRAM_CONTEXT_TIMEOUT_SECONDS = 1.5 _CONSUMER_STATEMENT_TIMEOUT_MS = 1500 +_CONSUMER_READBACK_CACHE_TTL_SECONDS = 20.0 +_consumer_readback_cache: dict[ + tuple[int, int, str], tuple[float, dict[str, Any]] +] = {} +_consumer_readback_inflight: dict[ + tuple[int, int, str], asyncio.Task[dict[str, Any]] +] = {} logger = get_logger(__name__) @@ -58,7 +68,48 @@ async def load_latest_ai_agent_log_controlled_writeback_consumer_readback( *, project_id: str = DEFAULT_PROJECT_ID, ) -> dict[str, Any]: - """Return live consumer bindings for LOG controlled writeback receipts.""" + """Return live consumer bindings without duplicating the same DB read burst.""" + + loop = asyncio.get_running_loop() + cache_key = (id(loop), id(get_db_context), project_id) + cached = _consumer_readback_cache_get(cache_key=cache_key) + if cached is not None: + return cached + + task = _consumer_readback_inflight.get(cache_key) + coalesced = task is not None + if task is None: + task = loop.create_task( + _load_latest_ai_agent_log_controlled_writeback_consumer_readback_uncached( + project_id=project_id + ) + ) + _consumer_readback_inflight[cache_key] = task + + try: + payload = await asyncio.shield(task) + payload = _with_consumer_readback_freshness( + payload, + cache_hit=False, + cache_age_seconds=0.0, + inflight_coalesced=coalesced, + ) + if _consumer_readback_cacheable(payload): + _consumer_readback_cache[cache_key] = ( + time.monotonic(), + copy.deepcopy(payload), + ) + return payload + finally: + if task.done() and _consumer_readback_inflight.get(cache_key) is task: + _consumer_readback_inflight.pop(cache_key, None) + + +async def _load_latest_ai_agent_log_controlled_writeback_consumer_readback_uncached( + *, + project_id: str, +) -> dict[str, Any]: + """Read current receipts once; the public loader owns cache/coalescing.""" try: rows, consumer_rows = await _load_consumer_rows_with_retry(project_id=project_id) @@ -94,6 +145,68 @@ async def load_latest_ai_agent_log_controlled_writeback_consumer_readback( ) +def _consumer_readback_cache_get( + *, + cache_key: tuple[int, int, str], +) -> dict[str, Any] | None: + now = time.monotonic() + for key, (stored_at, _payload) in list(_consumer_readback_cache.items()): + if now - stored_at > _CONSUMER_READBACK_CACHE_TTL_SECONDS: + _consumer_readback_cache.pop(key, None) + + cached = _consumer_readback_cache.get(cache_key) + if cached is None: + return None + stored_at, payload = cached + return _with_consumer_readback_freshness( + payload, + cache_hit=True, + cache_age_seconds=max(0.0, now - stored_at), + inflight_coalesced=False, + ) + + +def _consumer_readback_cacheable(payload: Mapping[str, Any]) -> bool: + boundaries = payload.get("operation_boundaries") + return bool( + isinstance(boundaries, Mapping) + and boundaries.get("metadata_ledger_read_performed") is True + and boundaries.get("consumer_apply_receipt_read_performed") is True + ) + + +def _with_consumer_readback_freshness( + payload: Mapping[str, Any], + *, + cache_hit: bool, + cache_age_seconds: float, + inflight_coalesced: bool, +) -> dict[str, Any]: + result = copy.deepcopy(dict(payload)) + prior = result.get("readback_freshness") + observed_at = ( + str(prior.get("observed_at") or "") + if isinstance(prior, Mapping) + else "" + ) + result["readback_freshness"] = { + "observed_at": observed_at or datetime.now(UTC).isoformat(), + "cache_ttl_seconds": _CONSUMER_READBACK_CACHE_TTL_SECONDS, + "cache_hit": cache_hit, + "cache_age_seconds": round(cache_age_seconds, 3), + "stale": False, + "inflight_coalesced": inflight_coalesced, + } + return result + + +def _reset_consumer_readback_cache_for_tests() -> None: + """Reset process-local readback state between focused unit tests.""" + + _consumer_readback_cache.clear() + _consumer_readback_inflight.clear() + + async def _load_consumer_rows_with_retry( *, project_id: str, diff --git a/apps/api/tests/test_ai_agent_log_controlled_writeback_consumer_readback_api.py b/apps/api/tests/test_ai_agent_log_controlled_writeback_consumer_readback_api.py index 5b40b1d75..044c19d59 100644 --- a/apps/api/tests/test_ai_agent_log_controlled_writeback_consumer_readback_api.py +++ b/apps/api/tests/test_ai_agent_log_controlled_writeback_consumer_readback_api.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json import pytest @@ -63,6 +64,8 @@ class _FailingContext: @pytest.fixture(autouse=True) def _stub_telegram_alert_card_readback(monkeypatch): + consumer_module._reset_consumer_readback_cache_for_tests() + async def fake_alert_card_readback(**_kwargs): return _telegram_alert_card_readback() @@ -71,6 +74,8 @@ def _stub_telegram_alert_card_readback(monkeypatch): "list_ai_alert_card_delivery_readback", fake_alert_card_readback, ) + yield + consumer_module._reset_consumer_readback_cache_for_tests() def _ledger_rows() -> list[dict]: @@ -215,6 +220,69 @@ async def test_log_controlled_writeback_consumer_loader_reads_live_ledger(monkey ] +@pytest.mark.asyncio +async def test_log_controlled_writeback_consumer_loader_reuses_fresh_success( + monkeypatch, +): + fake_db = _FakeDb(_ledger_rows(), _consumer_receipt_rows()) + monkeypatch.setattr( + consumer_module, + "get_db_context", + lambda project_id: _FakeContext(fake_db), + ) + + first = await load_latest_ai_agent_log_controlled_writeback_consumer_readback() + second = await load_latest_ai_agent_log_controlled_writeback_consumer_readback() + + assert first["readback_freshness"]["cache_hit"] is False + assert second["readback_freshness"]["cache_hit"] is True + assert second["readback_freshness"]["stale"] is False + assert second["readback_freshness"]["cache_age_seconds"] >= 0 + assert fake_db.params == [ + {"operation_type": "log_controlled_writeback_dispatched"}, + {"operation_type": "log_controlled_writeback_consumed"}, + ] + + +@pytest.mark.asyncio +async def test_log_controlled_writeback_consumer_loader_coalesces_concurrent_reads( + monkeypatch, +): + started = asyncio.Event() + release = asyncio.Event() + call_count = 0 + + async def delayed_rows(*, project_id: str): + nonlocal call_count + assert project_id == "awoooi" + call_count += 1 + started.set() + await release.wait() + return _ledger_rows(), _consumer_receipt_rows() + + monkeypatch.setattr( + consumer_module, + "_load_consumer_rows_with_retry", + delayed_rows, + ) + + first_task = asyncio.create_task( + load_latest_ai_agent_log_controlled_writeback_consumer_readback() + ) + await started.wait() + second_task = asyncio.create_task( + load_latest_ai_agent_log_controlled_writeback_consumer_readback() + ) + await asyncio.sleep(0) + release.set() + first, second = await asyncio.gather(first_task, second_task) + + assert call_count == 1 + assert first["readback_freshness"]["inflight_coalesced"] is False + assert second["readback_freshness"]["inflight_coalesced"] is True + assert first["consumer_bindings"] == second["consumer_bindings"] + + @pytest.mark.asyncio async def test_log_controlled_writeback_consumer_loader_reads_apply_receipts(monkeypatch): fake_db = _FakeDb(_ledger_rows(), _consumer_receipt_rows()) @@ -270,6 +338,39 @@ async def test_log_controlled_writeback_consumer_loader_degrades_on_db_pressure( assert payload["operation_boundaries"]["runtime_target_write_performed"] is False +@pytest.mark.asyncio +async def test_log_controlled_writeback_consumer_loader_does_not_cache_db_fallback( + monkeypatch, +): + monkeypatch.setattr( + consumer_module, + "get_db_context", + lambda project_id: _FailingContext(), + ) + direct_attempt_count = 0 + + async def fake_direct_unavailable(*, project_id: str): + nonlocal direct_attempt_count + assert project_id == "awoooi" + direct_attempt_count += 1 + return None + + monkeypatch.setattr( + consumer_module, + "_load_consumer_rows_direct", + fake_direct_unavailable, + ) + + first = await load_latest_ai_agent_log_controlled_writeback_consumer_readback() + second = await load_latest_ai_agent_log_controlled_writeback_consumer_readback() + + assert direct_attempt_count == 2 + assert first["readback_freshness"]["cache_hit"] is False + assert second["readback_freshness"]["cache_hit"] is False + assert first["readback"]["db_read_status"] == "unavailable" + assert second["readback"]["db_read_status"] == "unavailable" + + @pytest.mark.asyncio async def test_log_controlled_writeback_consumer_loader_uses_direct_fallback( monkeypatch,