fix(api): coalesce consumer readback DB queries
Some checks failed
CD Pipeline / workflow-shape (push) Successful in 0s
CD Pipeline / cancel-stale-cd (push) Has been skipped
CD Pipeline / tests (push) Successful in 1m20s
CD Pipeline / post-deploy-checks (push) Has been cancelled
CD Pipeline / build-and-deploy (push) Has been cancelled
Some checks failed
CD Pipeline / workflow-shape (push) Successful in 0s
CD Pipeline / cancel-stale-cd (push) Has been skipped
CD Pipeline / tests (push) Successful in 1m20s
CD Pipeline / post-deploy-checks (push) Has been cancelled
CD Pipeline / build-and-deploy (push) Has been cancelled
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user