fix(embeddings): the index refresh loses the race it used to deadlock (#3262)
CI & Build / Python lint (push) Successful in 5s
CI & Build / Plugin hooks (push) Successful in 16s
CI & Build / TypeScript typecheck (push) Successful in 37s
CI & Build / integration (push) Successful in 38s
CI & Build / Python tests (push) Failing after 55s
CI & Build / Build & push image (push) Skipped
CI & Build / Python lint (push) Successful in 5s
CI & Build / Plugin hooks (push) Successful in 16s
CI & Build / TypeScript typecheck (push) Successful in 37s
CI & Build / integration (push) Successful in 38s
CI & Build / Python tests (push) Failing after 55s
CI & Build / Build & push image (push) Skipped
An embedding refresh replaces a record's vectors as delete-then-insert, which takes the chunk rows first and the parent row second (via the insert's foreign key). A cascading delete of the parent takes exactly those two locks in the other order. Postgres calls the cycle a deadlock and kills one side: sometimes the detached embedder, silently, and sometimes the user's delete, as a 500 on an operation that should have worked. Both upserts now claim the parent row with FOR KEY SHARE NOWAIT before touching any chunk row. That removes the cycle instead of narrowing it — either the embedder is first and the delete queues behind it, or the delete already holds the row and the embedder loses at once, which is the side designed to lose. FOR KEY SHARE is the lock the insert would take anyway, so an ordinary edit is unaffected. The note twin, recorded as unverified on the issue, has the same shape and the same fix; a trash purge is the hard delete that reaches it. Unit tests pin the ORDER and the lock mode by compiling the statement; the integration pair holds a real delete open in one transaction and proves the embedder returns having written nothing, with a deadline so a regression fails instead of hanging.
This commit is contained in:
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+6
-2
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user