fix(crud): preserve cache invalidation across get_or_create retry
`get_or_create_peers` and `get_or_create_scopes` mutate existing rows, then insert new ones inside `db.begin_nested()`. When a concurrent writer creates one of those rows first, the insert raises IntegrityError and the function retries. `begin_nested()` autoflushes the pending UPDATEs *before* opening the savepoint, so the rollback neither undoes them nor expires the now-clean ORM state. The retry then compared already-updated values, found no change, and dropped those peers from `changed_peers` — skipping the cache purge while the row change committed anyway, leaving entries stale until the 300s TTL. Carry the mutated names into the retry via `_pending_invalidation` so the purge cannot be lost. The scopes facade mirrors `get_or_create_peers`, so both copies carried this. The peer path is pre-existing and runs on every message ingest. Also add the missing /v3 prefix to `_SCOPES_ROUTE_GUIDANCE`, which pointed callers at a 404. Adds tests/crud/test_get_or_create_retry_invalidation.py, which drives a real racing session and fails without this change.
This commit is contained in:
Vineeth Voruganti
parent
a17ce4176f
commit
679aa25388
|
|
@ -43,6 +43,7 @@ async def get_or_create_peers(
|
|||
peers: list[schemas.PeerCreate],
|
||||
*,
|
||||
_retry: bool = False,
|
||||
_pending_invalidation: list[str] | None = None,
|
||||
) -> GetOrCreateResult[list[models.Peer]]:
|
||||
"""
|
||||
Get an existing list of peers or create new peers if they don't exist.
|
||||
|
|
@ -53,6 +54,8 @@ async def get_or_create_peers(
|
|||
workspace_name: Name of the workspace
|
||||
peers: List of peer creation schemas
|
||||
_retry: Whether to retry the operation
|
||||
_pending_invalidation: Names of peers already mutated by a prior attempt,
|
||||
whose cache keys must still be purged. See the retry branch below.
|
||||
|
||||
Returns:
|
||||
GetOrCreateResult containing the list of peers and whether any were created
|
||||
|
|
@ -123,11 +126,26 @@ async def get_or_create_peers(
|
|||
raise ConflictException(
|
||||
f"Unable to create or get peers: {peer_names}"
|
||||
) from None
|
||||
return await get_or_create_peers(db, workspace_name, peers, _retry=True)
|
||||
# `begin_nested()` autoflushes the mutations above *before* opening the
|
||||
# savepoint, so they are already committed-in-transaction and the rollback
|
||||
# doesn't undo them — nor does it expire the now-clean ORM state. The retry
|
||||
# would therefore compare already-updated values, find no change, and skip
|
||||
# the purge. Carry the names forward so the invalidation can't be lost.
|
||||
return await get_or_create_peers(
|
||||
db,
|
||||
workspace_name,
|
||||
peers,
|
||||
_retry=True,
|
||||
_pending_invalidation=(_pending_invalidation or [])
|
||||
+ [p.name for p in changed_peers],
|
||||
)
|
||||
|
||||
# Capture peer names eagerly so the closure holds plain strings, not ORM objects
|
||||
_cache_keys_to_invalidate = [
|
||||
peer_cache_key(workspace_name, p.name) for p in changed_peers + new_peers
|
||||
peer_cache_key(workspace_name, name)
|
||||
for name in dict.fromkeys(
|
||||
(_pending_invalidation or []) + [p.name for p in changed_peers + new_peers]
|
||||
)
|
||||
]
|
||||
|
||||
async def _invalidate_peer_cache():
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ async def get_or_create_scopes(
|
|||
scopes: list[schemas.ScopeCreate],
|
||||
*,
|
||||
_retry: bool = False,
|
||||
_pending_invalidation: list[str] | None = None,
|
||||
) -> GetOrCreateResult[list[models.Peer]]:
|
||||
"""
|
||||
Get existing scopes or create new ones if they don't exist.
|
||||
|
|
@ -66,6 +67,9 @@ async def get_or_create_scopes(
|
|||
db: Database session
|
||||
workspace_name: Name of the workspace
|
||||
scopes: List of scope creation schemas (unprefixed names)
|
||||
_retry: Whether this is the retry attempt
|
||||
_pending_invalidation: Names of scope peers already mutated by a prior
|
||||
attempt, whose cache keys must still be purged. See the retry branch.
|
||||
|
||||
Returns:
|
||||
GetOrCreateResult containing the backing peers and whether any were
|
||||
|
|
@ -122,10 +126,24 @@ async def get_or_create_scopes(
|
|||
raise ConflictException(
|
||||
f"Unable to create or get scopes: {sorted(peer_names)}"
|
||||
) from None
|
||||
return await get_or_create_scopes(db, workspace_name, scopes, _retry=True)
|
||||
# `begin_nested()` autoflushes the mutations above *before* opening the
|
||||
# savepoint, so they survive the rollback and leave the ORM state clean —
|
||||
# the retry would compare already-updated values, find no change, and skip
|
||||
# the purge. Carry the names forward so the invalidation can't be lost.
|
||||
return await get_or_create_scopes(
|
||||
db,
|
||||
workspace_name,
|
||||
scopes,
|
||||
_retry=True,
|
||||
_pending_invalidation=(_pending_invalidation or [])
|
||||
+ [p.name for p in changed_peers],
|
||||
)
|
||||
|
||||
_cache_keys_to_invalidate = [
|
||||
peer_cache_key(workspace_name, p.name) for p in changed_peers + new_peers
|
||||
peer_cache_key(workspace_name, name)
|
||||
for name in dict.fromkeys(
|
||||
(_pending_invalidation or []) + [p.name for p in changed_peers + new_peers]
|
||||
)
|
||||
]
|
||||
|
||||
async def _invalidate_peer_cache():
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ router = APIRouter(
|
|||
# scopes facade so the observer mechanics stay internal.
|
||||
_SCOPES_ROUTE_GUIDANCE = (
|
||||
"Scope membership is managed via the scopes routes "
|
||||
"(/workspaces/{workspace_id}/scopes/{scope_id}/sessions) or the `scopes` "
|
||||
"(/v3/workspaces/{workspace_id}/scopes/{scope_id}/sessions) or the `scopes` "
|
||||
"field at session creation."
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -207,10 +207,10 @@ def _create_store_by_type(store_type: str) -> VectorStore:
|
|||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"VECTOR_STORE.TYPE is set to 'lancedb', but the 'lancedb' package "
|
||||
"is not installed (for example on macOS Intel, where it is omitted "
|
||||
"from dependencies because PyPI has no wheel). "
|
||||
"Use TYPE 'pgvector' or 'turbopuffer', or install lancedb manually. "
|
||||
f"Original import error: {exc}"
|
||||
+ "is not installed (for example on macOS Intel, where it is omitted "
|
||||
+ "from dependencies because PyPI has no wheel). "
|
||||
+ "Use TYPE 'pgvector' or 'turbopuffer', or install lancedb manually. "
|
||||
+ f"Original import error: {exc}"
|
||||
) from exc
|
||||
|
||||
return LanceDBVectorStore()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,196 @@
|
|||
"""Regression tests for cache invalidation across the get_or_create retry path.
|
||||
|
||||
`get_or_create_peers` / `get_or_create_scopes` mutate existing rows, then insert
|
||||
new ones inside `db.begin_nested()`. A concurrent writer that creates one of those
|
||||
rows first makes the insert raise `IntegrityError`, and the function retries.
|
||||
|
||||
The subtlety: `begin_nested()` autoflushes the pending mutations *before* opening
|
||||
the savepoint, so the rollback neither undoes them nor expires the ORM state. A
|
||||
retry that recomputed "what changed" from that state would see no change and skip
|
||||
the cache purge — while the row change still commits anyway, leaving the cache
|
||||
stale until TTL. These tests pin the purge.
|
||||
|
||||
The race is real (a second session committing a real row, producing a real
|
||||
IntegrityError from the database); only its *timing* is made deterministic, by
|
||||
hooking the one point that sits between the SELECT and the INSERT.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
from sqlalchemy.ext.asyncio import (
|
||||
AsyncEngine,
|
||||
AsyncSession,
|
||||
AsyncSessionTransaction,
|
||||
async_sessionmaker,
|
||||
)
|
||||
|
||||
from src import crud, models, schemas
|
||||
from src.crud.peer import peer_cache_key
|
||||
from src.crud.scope import SCOPE_PEER_CONFIGURATION
|
||||
from src.utils.scopes import scope_peer_name
|
||||
|
||||
|
||||
class _RaceOnBeginNested:
|
||||
"""Commit a racing row on entry to `begin_nested()`, then delegate.
|
||||
|
||||
That entry point is after the function's SELECT and metadata mutation but
|
||||
before its INSERT flushes — precisely the window a real concurrent writer
|
||||
has to slip through to trigger the IntegrityError retry.
|
||||
"""
|
||||
|
||||
_db: AsyncSession
|
||||
_engine: AsyncEngine
|
||||
_rows: list[models.Peer]
|
||||
_real: AsyncSessionTransaction | None
|
||||
fired: bool
|
||||
|
||||
def __init__(self, db: AsyncSession, engine: AsyncEngine, rows: list[models.Peer]):
|
||||
self._db = db
|
||||
self._engine = engine
|
||||
self._rows = rows
|
||||
self._real = None
|
||||
self.fired = False
|
||||
|
||||
def __call__(self):
|
||||
return self
|
||||
|
||||
async def __aenter__(self):
|
||||
if self._rows:
|
||||
Session = async_sessionmaker(bind=self._engine, expire_on_commit=False)
|
||||
async with Session() as other:
|
||||
other.add_all(self._rows)
|
||||
await other.commit()
|
||||
self._rows = [] # race only once; the retry must succeed
|
||||
self.fired = True
|
||||
self._real = AsyncSession.begin_nested(self._db)
|
||||
return await self._real.__aenter__()
|
||||
|
||||
async def __aexit__(self, *exc_info: object):
|
||||
assert self._real is not None
|
||||
return await self._real.__aexit__(*exc_info)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_peer_retry_still_invalidates_mutated_peer(
|
||||
db_session: AsyncSession,
|
||||
db_engine: AsyncEngine,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A peer mutated before a losing race still gets its cache key purged."""
|
||||
test_workspace, existing_peer = sample_data
|
||||
racer_name = str(generate_nanoid())
|
||||
|
||||
# Give the existing peer metadata we will then change, so it is a real update.
|
||||
existing_peer.h_metadata = {"v": "old"}
|
||||
await db_session.commit()
|
||||
|
||||
race = _RaceOnBeginNested(
|
||||
db_session,
|
||||
db_engine,
|
||||
[models.Peer(name=racer_name, workspace_name=test_workspace.name)],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("src.crud.peer.safe_cache_delete", new=AsyncMock()) as mock_delete,
|
||||
patch.object(db_session, "begin_nested", race),
|
||||
):
|
||||
result = await crud.get_or_create_peers(
|
||||
db_session,
|
||||
test_workspace.name,
|
||||
[
|
||||
schemas.PeerCreate(name=existing_peer.name, metadata={"v": "new"}),
|
||||
schemas.PeerCreate(name=racer_name),
|
||||
],
|
||||
)
|
||||
await db_session.commit()
|
||||
await result.post_commit()
|
||||
|
||||
assert race.fired, "the race must actually have fired"
|
||||
|
||||
purged = {call.args[0] for call in mock_delete.await_args_list}
|
||||
assert (
|
||||
peer_cache_key(test_workspace.name, existing_peer.name) in purged
|
||||
), "the mutated peer's cache key must still be purged after the retry"
|
||||
|
||||
# The mutation really did land — which is what makes a missed purge stale.
|
||||
await db_session.refresh(existing_peer)
|
||||
assert existing_peer.h_metadata == {"v": "new"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scope_retry_still_invalidates_mutated_scope(
|
||||
db_session: AsyncSession,
|
||||
db_engine: AsyncEngine,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""Same guarantee for the scopes facade, which mirrors get_or_create_peers."""
|
||||
test_workspace, _ = sample_data
|
||||
kept_scope, racing_scope = str(generate_nanoid()), str(generate_nanoid())
|
||||
|
||||
seeded = await crud.get_or_create_scopes(
|
||||
db_session,
|
||||
test_workspace.name,
|
||||
[schemas.ScopeCreate(name=kept_scope, metadata={"v": "old"})],
|
||||
)
|
||||
await db_session.commit()
|
||||
await seeded.post_commit()
|
||||
|
||||
# The racer creates the second scope's backing peer — as a *valid* scope peer,
|
||||
# so the flow reaches the insert rather than tripping the legacy-collision 409.
|
||||
race = _RaceOnBeginNested(
|
||||
db_session,
|
||||
db_engine,
|
||||
[
|
||||
models.Peer(
|
||||
name=scope_peer_name(racing_scope),
|
||||
workspace_name=test_workspace.name,
|
||||
configuration=dict(SCOPE_PEER_CONFIGURATION),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("src.crud.scope.safe_cache_delete", new=AsyncMock()) as mock_delete,
|
||||
patch.object(db_session, "begin_nested", race),
|
||||
):
|
||||
result = await crud.get_or_create_scopes(
|
||||
db_session,
|
||||
test_workspace.name,
|
||||
[
|
||||
schemas.ScopeCreate(name=kept_scope, metadata={"v": "new"}),
|
||||
schemas.ScopeCreate(name=racing_scope),
|
||||
],
|
||||
)
|
||||
await db_session.commit()
|
||||
await result.post_commit()
|
||||
|
||||
assert race.fired, "the race must actually have fired"
|
||||
|
||||
purged = {call.args[0] for call in mock_delete.await_args_list}
|
||||
assert (
|
||||
peer_cache_key(test_workspace.name, scope_peer_name(kept_scope)) in purged
|
||||
), "the mutated scope peer's cache key must still be purged after the retry"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_peer_no_race_does_not_invalidate_unchanged_peer(
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""Baseline: with no race, an unchanged peer is not purged."""
|
||||
test_workspace, existing_peer = sample_data
|
||||
existing_peer.h_metadata = {"v": "same"}
|
||||
await db_session.commit()
|
||||
|
||||
with patch("src.crud.peer.safe_cache_delete", new=AsyncMock()) as mock_delete:
|
||||
result = await crud.get_or_create_peers(
|
||||
db_session,
|
||||
test_workspace.name,
|
||||
[schemas.PeerCreate(name=existing_peer.name, metadata={"v": "same"})],
|
||||
)
|
||||
await db_session.commit()
|
||||
await result.post_commit()
|
||||
|
||||
assert mock_delete.await_count == 0, "an unchanged peer must not be purged"
|
||||
|
|
@ -2,7 +2,9 @@
|
|||
|
||||
import pytest
|
||||
|
||||
from src.cache.client import _redact_cache_url
|
||||
from src.cache.client import (
|
||||
_redact_cache_url, # pyright: ignore[reportPrivateUsage]
|
||||
)
|
||||
|
||||
|
||||
class TestRedactCacheUrl:
|
||||
|
|
|
|||
Loading…
Reference in New Issue