diff --git a/src/scribe/services/embeddings.py b/src/scribe/services/embeddings.py index fb8badd..a905b1e 100644 --- a/src/scribe/services/embeddings.py +++ b/src/scribe/services/embeddings.py @@ -344,6 +344,49 @@ def chunk_document(title: str | None, body: str | None) -> list[str]: return chunks +async def _claim_parent_row(session, id_column, row_id: int, label: str) -> bool: + """Lock the record a vector belongs to BEFORE rewriting that vector (#3262). + + An embedding write and a cascading delete of the same record take the same + two row locks in OPPOSITE orders. The embedder deletes the old chunk rows + and then, on INSERT, needs the foreign key's lock on the parent; a delete + of the parent — or of the rulebook, topic or project above it — locks the + parent first and cascades down into the chunk rows. That is a cycle, and + Postgres breaks it by killing one side at random: sometimes the embedding + write, which is swallowed and invisible, and sometimes the operator's + delete, which surfaces as a 500 on an operation that should have worked. + + Claiming the parent first REMOVES the cycle rather than narrowing it. + Either the embedder arrives first and the delete waits its turn behind it, + or the delete already holds the row and NOWAIT makes the embedder lose at + once. The embedder is the side that should lose: a skipped refresh costs a + stale vector until the next write or the startup backfill, and the other + outcome costs a person their request. + + FOR KEY SHARE, not FOR UPDATE — it is precisely the lock the INSERT's + foreign key would take anyway, so it conflicts with a delete of the parent + and with nothing else. An ordinary edit of the same record, or a second + refresh racing this one, is unaffected. + + Returns False when the row is locked or already gone; the caller skips. + """ + try: + held = (await session.execute( + select(id_column) + .where(id_column == row_id) + .with_for_update(key_share=True, nowait=True) + )).scalar_one_or_none() + except Exception: + # LockNotAvailable: this record is being deleted right now. Not an + # error — the delete wins by design. + logger.debug("Skipping embedding for %s %d — row is being deleted", label, row_id) + return False + if held is None: + logger.debug("Skipping embedding for %s %d — row is gone", label, row_id) + return False + return True + + async def upsert_note_embedding( note_id: int, user_id: int, title: str | None, body: str | None ) -> None: @@ -380,6 +423,8 @@ async def upsert_note_embedding( try: async with async_session() as session: + if not await _claim_parent_row(session, Note.id, note_id, "note"): + return await session.execute( delete(NoteEmbedding).where(NoteEmbedding.note_id == note_id) ) @@ -666,6 +711,8 @@ async def upsert_rule_embedding( replacement is atomic per rule so a concurrent read sees the old chunk set or the new one, never a mixture. """ + from scribe.models.rulebook import Rule # runtime import: see TYPE_CHECKING above + doc_title, doc_body = rule_document(title, statement, when_to_apply) chunks = chunk_document(doc_title, doc_body) try: @@ -688,6 +735,8 @@ async def upsert_rule_embedding( try: async with async_session() as session: + if not await _claim_parent_row(session, Rule.id, rule_id, "rule"): + return await session.execute( delete(RuleEmbedding).where(RuleEmbedding.rule_id == rule_id) ) diff --git a/src/scribe/services/notes.py b/src/scribe/services/notes.py index db8d61d..4deee9c 100644 --- a/src/scribe/services/notes.py +++ b/src/scribe/services/notes.py @@ -76,6 +76,10 @@ def embed_note(note) -> None: exceptions are swallowed because a record that saved must not fail on its index refresh. No running loop (unit tests, scripts) is an ordinary case, not an error. + + Detaching also means this task races anything that deletes the note out + from under it. That is not handled here: `upsert_note_embedding` claims + the note's row before touching its vectors, and loses if it can't (#3262). """ try: import asyncio diff --git a/src/scribe/services/rulebooks.py b/src/scribe/services/rulebooks.py index d2378c1..ba11b69 100644 --- a/src/scribe/services/rulebooks.py +++ b/src/scribe/services/rulebooks.py @@ -379,6 +379,10 @@ def _refresh_rule_embedding(rule: Rule) -> None: swallowed because a rule that SAVED must not fail on its index refresh — a stale vector costs a missed search hit, a raised exception costs the write. No running loop (unit tests, scripts) is ordinary, not an error. + + Detaching also means this task races anything that deletes the rule out + from under it. That is not handled here: `upsert_rule_embedding` claims + the rule's row before touching its vectors, and loses if it can't (#3262). """ try: import asyncio diff --git a/tests/helpers.py b/tests/helpers.py index 45ac4a0..c23516c 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -59,14 +59,18 @@ def tool_doc(module: str, name: str) -> str: return _re.sub(r"\s+", " ", fn.__doc__) -def compiled_sql(element) -> str: +def compiled_sql(element, dialect=None) -> str: """A SQLAlchemy clause or statement rendered as literal SQL text. For asserting on the shape of a predicate without a database — which is how the visibility clauses and the knowledge facets are both tested. Was a private copy in each of those modules before #3128 needed a third. + + Pass `dialect` when the assertion is about something only one backend + renders — a Postgres row-lock mode, say. The generic dialect is enough for + a predicate's shape and would quietly drop the rest. """ - return str(element.compile(compile_kwargs={"literal_binds": True})) + return str(element.compile(dialect=dialect, compile_kwargs={"literal_binds": True})) def make_mock_session() -> AsyncMock: diff --git a/tests/test_embedding_yields_to_a_delete.py b/tests/test_embedding_yields_to_a_delete.py new file mode 100644 index 0000000..ba50388 --- /dev/null +++ b/tests/test_embedding_yields_to_a_delete.py @@ -0,0 +1,118 @@ +"""The embedding refresh must LOSE to a delete, not race it (#3262). + +Both upserts replace a record's vectors as delete-then-insert. That takes two +row locks — the chunk rows, then the parent row via the insert's foreign key — +in the exact reverse of the order a cascading delete of the parent takes them. +Postgres calls that a deadlock and kills one side at random, which sometimes +means killing the user's delete. + +These pin the ORDER, not the outcome: the claim on the parent goes first, and +when the claim fails nothing else in the transaction runs. Compiling the +statement is the only way to assert on a lock mode without a database — the +integration twin (test_integration_embedding_yields_to_delete.py) proves the +behaviour against a real one. +""" +from unittest.mock import AsyncMock, MagicMock, patch + +from sqlalchemy.dialects import postgresql +from sqlalchemy.exc import OperationalError + +from scribe.services import embeddings as emb +from tests.helpers import compiled_sql + +ONE_VECTOR = [[0.0] * 384] + +# The lock mode is a Postgres extension — the generic dialect renders a plain +# FOR UPDATE and would pass an assertion that proves nothing. +PG = postgresql.dialect() + + +def _mock_session(lock_result: object = 7, execute_side_effect=None): + """A session stand-in whose first execute answers the parent-row claim.""" + session = MagicMock() + claimed = MagicMock() + claimed.scalar_one_or_none.return_value = lock_result + if execute_side_effect is not None: + session.execute = AsyncMock(side_effect=execute_side_effect) + else: + session.execute = AsyncMock(return_value=claimed) + session.commit = AsyncMock() + session.add = MagicMock() + ctx = MagicMock() + ctx.__aenter__ = AsyncMock(return_value=session) + ctx.__aexit__ = AsyncMock(return_value=False) + return session, ctx + + +def _lock_unavailable() -> OperationalError: + """What asyncpg raises through SQLAlchemy when NOWAIT can't take the row.""" + return OperationalError("SELECT ...", {}, Exception("lock not available")) + + +async def test_a_note_refresh_claims_the_row_before_rewriting_its_vectors(): + """The claim is FIRST, and it is FOR KEY SHARE NOWAIT. + + FOR KEY SHARE because that is exactly the lock the insert's foreign key + takes anyway — it conflicts with a delete of the note and with nothing + else, so an ordinary edit is unaffected. NOWAIT because the whole point is + to lose immediately rather than queue up behind the delete and hold the + chunk rows while doing it. + """ + session, ctx = _mock_session() + with ( + patch.object(emb, "async_session", return_value=ctx), + patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)), + ): + await emb.upsert_note_embedding(7, 42, "T", "a short body") + + claim, replace = [c.args[0] for c in session.execute.call_args_list][:2] + assert compiled_sql(claim, dialect=PG).startswith("SELECT notes.id") + assert "FOR KEY SHARE NOWAIT" in compiled_sql(claim, dialect=PG) + assert compiled_sql(replace, dialect=PG).startswith("DELETE FROM note_embeddings") + session.add.assert_called() + + +async def test_a_rule_refresh_claims_the_row_before_rewriting_its_vectors(): + """The rule twin — the path the reported deadlock actually took.""" + session, ctx = _mock_session() + with ( + patch.object(emb, "async_session", return_value=ctx), + patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)), + ): + await emb.upsert_rule_embedding(9, "T", "a short statement", "on write") + + claim, replace = [c.args[0] for c in session.execute.call_args_list][:2] + assert compiled_sql(claim, dialect=PG).startswith("SELECT rules.id") + assert "FOR KEY SHARE NOWAIT" in compiled_sql(claim, dialect=PG) + assert compiled_sql(replace, dialect=PG).startswith("DELETE FROM rule_embeddings") + session.add.assert_called() + + +async def test_a_record_being_deleted_is_left_alone_rather_than_raced(): + """The claim failing ends the write — it does not fall through to the + delete-and-insert that would take the locks in the losing order.""" + session, ctx = _mock_session(execute_side_effect=_lock_unavailable()) + with ( + patch.object(emb, "async_session", return_value=ctx), + patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)), + ): + await emb.upsert_rule_embedding(9, "T", "a short statement", "on write") + + assert session.execute.await_count == 1, "it stopped at the claim" + session.add.assert_not_called() + session.commit.assert_not_awaited() + + +async def test_a_record_already_gone_is_not_re_embedded(): + """A vector inserted for a row that no longer exists is either a foreign + key violation or, worse, a resurrected chunk. Nothing to refresh.""" + session, ctx = _mock_session(lock_result=None) + with ( + patch.object(emb, "async_session", return_value=ctx), + patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)), + ): + await emb.upsert_note_embedding(7, 42, "T", "a short body") + + assert session.execute.await_count == 1 + session.add.assert_not_called() + session.commit.assert_not_awaited() diff --git a/tests/test_integration_embedding_yields_to_delete.py b/tests/test_integration_embedding_yields_to_delete.py new file mode 100644 index 0000000..be17d73 --- /dev/null +++ b/tests/test_integration_embedding_yields_to_delete.py @@ -0,0 +1,173 @@ +"""#3262 against a real Postgres: the embedder loses the race, it doesn't run it. + +The reported failure was a deadlock — `DELETE FROM rulebooks` killed by the +server while a detached `upsert_rule_embedding` held the other half of the +cycle. It cannot be reproduced with mocks, because there is nothing to +deadlock: the whole bug lives in the ORDER two transactions take two row +locks, which only a lock manager can adjudicate. + +So each test here holds a real delete open in one transaction and calls the +embedder in another. What is being pinned is that the embedder RETURNS — +promptly, having written nothing. Before the fix it would sit on the chunk +rows waiting for a delete that is itself waiting on the insert's foreign key, +and the test would hang rather than fail, which is why every call carries a +deadline (rule 156). +""" +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest +import pytest_asyncio +from sqlalchemy import delete, select + +from scribe.models import async_session +from scribe.models.embedding import NoteEmbedding, RuleEmbedding +from scribe.models.note import Note +from scribe.models.rulebook import Rulebook +from scribe.services import embeddings as emb +from scribe.services import rulebooks as rulebooks_svc +from tests.helpers import ensure_user + +pytestmark = [pytest.mark.integration, pytest.mark.usefixtures("_dispose_engine")] + +OWNER_USERNAME = "embed_lock_owner" + +# Generous, because the assertion is "it did not block indefinitely", not "it +# was fast". A machine under load must not turn this into a flake; a genuinely +# blocked embedder never returns at all, so no honest run comes near this. +YIELD_DEADLINE_SECONDS = 20 + +# What the embedder would write if it wrongly went ahead. Distinct from the +# text the fixture's own create_* wrote, so the assertion cannot be satisfied +# by rows that were already there. +SENTINEL = "sentinelvector" + +ONE_VECTOR = [[0.0] * 384] + +# The fixture's own embedding task runs the REAL embedder, which either loads a +# model or gives up; both are bounded well inside this. +SETTLE_DEADLINE_SECONDS = 30 + + +@pytest_asyncio.fixture +async def seeded(): + """A rule and a note to race against. + + CLEANED AT SETUP, NOT TEARDOWN — the same constraint #3241 hit and the + reason this file exists. `create_rule` fires its own detached embedding + task; a teardown that deleted the rulebook would be racing exactly the + thing under test, on a loop that is closing. + """ + async with async_session() as s: + owner = await ensure_user(s, OWNER_USERNAME) + uid = owner.id + await s.commit() + for book in (await s.execute( + select(Rulebook).where(Rulebook.owner_user_id == uid) + )).scalars().all(): + await s.delete(book) + for note in (await s.execute( + select(Note).where(Note.user_id == uid) + )).scalars().all(): + await s.delete(note) + await s.commit() + + book = await rulebooks_svc.create_rulebook(uid, "Lock fixtures") + topic = await rulebooks_svc.create_topic(book.id, uid, "locks") + rule = await rulebooks_svc.create_rule( + topic.id, uid, "A rule with vectors", + "Something for the embedder to index.", + ) + async with async_session() as s: + note = Note(user_id=uid, title="A note with vectors", body="Body text.") + s.add(note) + await s.commit() + note_id = note.id + + await _settle_detached_writes() + return {"uid": uid, "book_id": book.id, "rule_id": rule.id, "note_id": note_id} + + +async def _settle_detached_writes() -> None: + """Let `create_rule`'s own fire-and-forget embedding task finish. + + It is the same detached write these tests are about, aimed at the same + rule, and left in flight it would land in the middle of an assertion about + that rule's rows. Bounded, and a timeout is not a failure — the tests below + carry their own deadlines, and this is only tidying the start line. + """ + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + if pending: + await asyncio.wait(pending, timeout=SETTLE_DEADLINE_SECONDS) + + +async def _rule_chunks(rule_id: int) -> list[str]: + async with async_session() as s: + return list((await s.execute( + select(RuleEmbedding.chunk_text).where(RuleEmbedding.rule_id == rule_id) + )).scalars().all()) + + +async def _note_chunks(note_id: int) -> list[str]: + async with async_session() as s: + return list((await s.execute( + select(NoteEmbedding.chunk_text).where(NoteEmbedding.note_id == note_id) + )).scalars().all()) + + +async def test_a_rule_refresh_yields_to_a_delete_cascading_from_its_rulebook(seeded): + """The reported case, exactly: the delete lands on the RULEBOOK and reaches + the rule through two cascades, which is why nothing on the rule's own write + path could have seen it coming.""" + async with async_session() as blocker: + # Uncommitted on purpose — the cascade's locks are held for as long as + # this transaction stays open, which is the state the embedder must + # decline to fight over. + await blocker.execute(delete(Rulebook).where(Rulebook.id == seeded["book_id"])) + try: + with patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)): + await asyncio.wait_for( + emb.upsert_rule_embedding( + seeded["rule_id"], SENTINEL, f"{SENTINEL} statement", + ), + timeout=YIELD_DEADLINE_SECONDS, + ) + finally: + await blocker.rollback() + + assert not any(SENTINEL in text for text in await _rule_chunks(seeded["rule_id"])), \ + "the embedder wrote into a rule that was being deleted" + + +async def test_a_note_refresh_yields_to_a_delete_of_the_note(seeded): + """The note twin, which #3262 recorded as unverified. Notes are soft-deleted + day to day, so the hard delete a trash purge issues is the one that can put + a lock on the row while a refresh is in flight.""" + async with async_session() as blocker: + await blocker.execute(delete(Note).where(Note.id == seeded["note_id"])) + try: + with patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)): + await asyncio.wait_for( + emb.upsert_note_embedding( + seeded["note_id"], seeded["uid"], SENTINEL, f"{SENTINEL} body", + ), + timeout=YIELD_DEADLINE_SECONDS, + ) + finally: + await blocker.rollback() + + assert not any(SENTINEL in text for text in await _note_chunks(seeded["note_id"])), \ + "the embedder wrote into a note that was being deleted" + + +async def test_an_uncontended_refresh_still_writes(seeded): + """The guard against the cheapest possible false pass: a claim that never + succeeds would satisfy both tests above while quietly ending semantic + search.""" + with patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)): + await emb.upsert_rule_embedding( + seeded["rule_id"], SENTINEL, f"{SENTINEL} statement", + ) + + assert any(SENTINEL in text for text in await _rule_chunks(seeded["rule_id"])), \ + "an unlocked rule was not embedded"