diff --git a/src/basic_memory/repository/accepted_note_search_repository.py b/src/basic_memory/repository/accepted_note_search_repository.py index 7e9ad07fa..e0ea349b2 100644 --- a/src/basic_memory/repository/accepted_note_search_repository.py +++ b/src/basic_memory/repository/accepted_note_search_repository.py @@ -15,6 +15,7 @@ delete_project_index_vector_rows, ) from basic_memory.repository.postgres_fts_chunks import split_postgres_fts_chunks +from basic_memory.repository.search_projection_lock import lock_entity_search_projection from basic_memory.repository.script_ngrams import build_script_ngrams type SearchIndexSqlValue = str | int | datetime | None @@ -177,6 +178,12 @@ async def refresh_entity( f"does not match repository project_id {self.project_id}" ) + # Trigger: the accepted write rewrites the same rows the indexer does. + # Why: two rewrites of one entity at once deadlock or hit search_index_pkey. + # Outcome: it takes the same per-entity turn, before its delete. + await lock_entity_search_projection( + session, project_id=row.project_id, entity_id=row.entity_id + ) await session.execute( DELETE_ACCEPTED_NOTE_SEARCH_SQL, {"entity_id": row.entity_id, "project_id": row.project_id}, @@ -224,6 +231,10 @@ async def delete_entity( entity_id: int, ) -> None: """Delete all accepted-note search rows for one entity.""" + # Same per-entity turn as every other rewrite of this entity's rows. + await lock_entity_search_projection( + session, project_id=self.project_id, entity_id=entity_id + ) await session.execute( DELETE_ACCEPTED_NOTE_SEARCH_SQL, {"entity_id": entity_id, "project_id": self.project_id}, diff --git a/src/basic_memory/repository/search_projection_lock.py b/src/basic_memory/repository/search_projection_lock.py new file mode 100644 index 000000000..4425904eb --- /dev/null +++ b/src/basic_memory/repository/search_projection_lock.py @@ -0,0 +1,49 @@ +"""Serialize rewrites of one entity's search projection on Postgres. + +Every refresh of an entity's search rows is delete-then-insert in one transaction: +delete every ``search_index`` row the entity owns, then insert the entity row, its +observation rows and its relation rows again. Two such transactions for the same +entity, running at once, collide in three ways under READ COMMITTED: + +- the second DELETE waits on the first's row locks, then re-checks rows the first + has already replaced, and the two wait on each other: ``deadlock detected``; +- the second DELETE's snapshot predates the first's INSERT, so it does not delete + the new rows, and the second INSERT then hits ``search_index_pkey`` on + ``(id, type, project_id)``. The upsert cannot absorb that: its conflict target is + the permalink index, and Postgres raises a violation on any other unique index; +- the same, on relation rows, whose ids are the relation table's. + +Saving one long note twice in quick succession is enough, because re-indexing a long +note takes longer than the gap between saves. Taking a transaction-scoped advisory +lock keyed on (project, entity) before the DELETE makes the rewrites of one entity +take turns, and the lock is released by the commit or rollback that ends the +transaction. Rewrites of different entities do not wait on each other. + +Lock order: a transaction must take this lock before it touches any of the entity's +``search_index`` rows. Every writer that rewrites an entity's projection does, so a +holder of the lock never waits on another writer of those rows. +""" + +from __future__ import annotations + +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession + +# The key namespaces the lock, so it cannot collide with an advisory lock some other +# component takes on a bare (int, int) pair. +_LOCK_SQL = text( + "SELECT pg_advisory_xact_lock(" + "hashtextextended('basic-memory:search-projection:' || :project_id || ':' || :entity_id, 0))" +) + + +async def lock_entity_search_projection( + session: AsyncSession, *, project_id: int, entity_id: int +) -> None: + """Wait until no other transaction is rewriting this entity's search rows. + + A no-op on SQLite, which already allows one writer at a time. + """ + if session.get_bind().dialect.name != "postgresql": + return + await session.execute(_LOCK_SQL, {"project_id": str(project_id), "entity_id": str(entity_id)}) diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index 559f203fa..83ecf6b48 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -24,6 +24,7 @@ SearchIndexRow, SearchRepository, ) +from basic_memory.repository.search_projection_lock import lock_entity_search_projection from basic_memory.repository.search_query import PreparedSearchQuery, relaxed_query_words from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.search_trace import SearchTraceCollector @@ -634,13 +635,6 @@ async def index_entity_data( f"permalink={entity.permalink} project_id={entity.project_id}" ) try: - replacement_content = content - if entity.is_markdown and replacement_content is None: - # Trigger: synchronized and legacy notes source search text from storage. - # Why: a transient read failure must preserve the last valid projection. - # Outcome: storage errors remain visible before any search rows are deleted. - replacement_content = await self.file_service.read_entity_content(entity) - # Trigger: every refresh replaces the entity's whole search projection. # Why: the delete used to commit on its own, so a timeout while writing the # replacement (Postgres FTS chunks under load, #1621) left an existing note @@ -649,6 +643,22 @@ async def index_entity_data( # Outcome: delete and replacement commit together; a failure rolls back to # the previous projection, or to none on a first index, and stays retryable. async with db.scoped_session(self.repository.session_maker) as session: + # Trigger: two refreshes of one entity run at once (two quick saves of + # a long note, or a save and its follow-up reindex). + # Why: on Postgres they deadlocked on the delete or hit search_index_pkey. + # Outcome: refreshes of one entity take turns; see search_projection_lock. + await lock_entity_search_projection( + session, project_id=entity.project_id, entity_id=entity.id + ) + replacement_content = content + if entity.is_markdown and replacement_content is None: + # Trigger: synchronized and legacy notes source search text from storage. + # Why: a transient read failure must preserve the last valid projection. + # Outcome: storage errors remain visible before any search rows are + # deleted, and the rollback keeps the previous projection. + # Read after taking the lock, so the last refresh to run indexes + # the file as it is now, not as it was when it started waiting. + replacement_content = await self.file_service.read_entity_content(entity) await self.repository.delete_by_entity_id(entity_id=entity.id, session=session) if entity.is_markdown: await self.index_entity_markdown(entity, replacement_content, session=session) diff --git a/tests/repository/test_accepted_note_search_repository.py b/tests/repository/test_accepted_note_search_repository.py index 4d3a00b48..f7a532590 100644 --- a/tests/repository/test_accepted_note_search_repository.py +++ b/tests/repository/test_accepted_note_search_repository.py @@ -59,6 +59,11 @@ async def test_refresh_entity_replaces_project_scoped_hot_search_row() -> None: await repository.refresh_entity(cast(AsyncSession, session), row) + # On Postgres the refresh first takes the entity's search projection lock, + # before it touches any row. + lock_sql, lock_params = session.executed.pop(0) + assert "pg_advisory_xact_lock" in lock_sql + assert lock_params == {"project_id": "7", "entity_id": "42"} assert len(session.executed) == 4 delete_sql, delete_params = session.executed[0] insert_sql, insert_params = session.executed[1] diff --git a/tests/services/test_search_refresh_concurrency.py b/tests/services/test_search_refresh_concurrency.py new file mode 100644 index 000000000..951ccd463 --- /dev/null +++ b/tests/services/test_search_refresh_concurrency.py @@ -0,0 +1,258 @@ +"""Overlapping refreshes of one entity's search rows all succeed. + +Saving a long note twice in quick succession ran two delete-then-insert refreshes of +the same entity at once. On Postgres the second one failed with ``deadlock detected`` +or with ``search_index_pkey`` on the entity row or a relation row, and the save that +triggered it answered 500. These tests run N refreshes of one long note at once. + +A real 80 KB note opens the race window by itself: indexing it takes longer than the +gap between saves. Here every refresh also pauses between its delete and its insert, +and the collision rounds start from an entity whose rows another writer has just +deleted, so the deletes have nothing to wait on and the inserts arrive together. The +negative controls remove the lock and show the same refreshes fail. +""" + +from __future__ import annotations + +import asyncio +import random +import string + +import pytest + +from basic_memory import db +from basic_memory.models import Entity +from basic_memory.repository.accepted_note_search_repository import AcceptedNoteSearchRepository +from basic_memory.repository.accepted_note_search_row import AcceptedNoteSearchRow +from basic_memory.services import search_service as search_service_module + +pytestmark = [pytest.mark.asyncio, pytest.mark.postgres] + +CONCURRENT_SAVES = 6 +NOTE_BYTES = 80_000 + + +def _long_body(version: int) -> str: + rnd = random.Random(version) + words = [ + "".join(rnd.choice(string.ascii_lowercase) for _ in range(rnd.randint(3, 10))) + for _ in range(NOTE_BYTES // 7) + ] + lines = [" ".join(words[i : i + 14]) for i in range(0, len(words), 14)] + return f"versionmarker{version}\n\n" + "\n\n".join(lines) + + +async def _reload(session_maker, entity_repository, entity: Entity) -> Entity: + async with db.scoped_session(session_maker) as session: + reloaded = await entity_repository.find_by_id(session, entity.id) + assert reloaded is not None + return reloaded + + +def _widen_the_window(monkeypatch, repository) -> None: + """Pause between a refresh's delete and its insert, while it holds the deleted rows.""" + original = repository.bulk_index_items + + async def pause_then_insert(search_index_rows, session=None): + await asyncio.sleep(0.05) + await original(search_index_rows, session) + + monkeypatch.setattr(repository, "bulk_index_items", pause_then_insert) + + +def _remove_the_lock(monkeypatch) -> None: + async def no_lock(session, *, project_id, entity_id): + return None + + monkeypatch.setattr(search_service_module, "lock_entity_search_projection", no_lock) + + +async def _projection(search_service, entity: Entity) -> list[tuple[str, int]]: + rows = await search_service.repository.get_entity_search_rows(entity.id) + return sorted((row.type, row.id) for row in rows if row.entity_id == entity.id) + + +def _skip_unless_postgres(db_backend) -> None: + if db_backend != "postgres": + pytest.skip("the collision is Postgres row locking; SQLite allows one writer") + + +async def _save_then_index(search_service, file_service, entity: Entity, version: int) -> None: + """What a save does: write the file, then refresh the index from storage.""" + path = file_service.get_entity_path(entity) + await asyncio.to_thread( + path.write_text, f"---\ntitle: {entity.title}\n---\n{_long_body(version)}\n" + ) + await search_service.index_entity_data(entity) + + +async def _save_round(search_service, file_service, entity: Entity, start: int) -> list[object]: + """N saves at once, from an entity whose old rows another writer just deleted. + + With no rows to wait on, every refresh's delete passes at once and all of them + reach their inserts together, which is when ``search_index_pkey`` fires. + """ + await search_service.repository.delete_by_entity_id(entity.id) + return await asyncio.gather( + *( + _save_then_index(search_service, file_service, entity, v) + for v in range(start, start + CONCURRENT_SAVES) + ), + return_exceptions=True, + ) + + +async def test_concurrent_saves_of_a_long_note_all_succeed_and_the_last_one_wins( + monkeypatch, + db_backend, + search_service, + file_service, + full_entity, + entity_repository, + session_maker, +): + _skip_unless_postgres(db_backend) + entity = await _reload(session_maker, entity_repository, full_entity) + await search_service.index_entity_data(entity, content="before") + before = await _projection(search_service, entity) + _widen_the_window(monkeypatch, search_service.repository) + + for round_start in (0, 100, 200): + results = await _save_round(search_service, file_service, entity, round_start) + assert [r for r in results if isinstance(r, BaseException)] == [] + + # The projection is whole and single: one row per entity, observation and + # relation, the same set as before the overlap. + assert await _projection(search_service, entity) == before + + # The index holds the text on disk now, which is the last save's. + on_disk = await file_service.read_entity_content(entity) + rows = await search_service.repository.get_entity_search_rows(entity.id) + [entity_row] = [row for row in rows if row.type == "entity"] + assert entity_row.content_snippet == on_disk + + +async def test_negative_control_without_the_lock_overlapping_saves_fail( + monkeypatch, + db_backend, + search_service, + file_service, + full_entity, + entity_repository, + session_maker, +): + _skip_unless_postgres(db_backend) + entity = await _reload(session_maker, entity_repository, full_entity) + await search_service.index_entity_data(entity, content="before") + _widen_the_window(monkeypatch, search_service.repository) + _remove_the_lock(monkeypatch) + + failures: list[BaseException] = [] + # The collision is a race, so a round can miss; ten rounds have not all missed. + for round_start in range(0, 1000, 100): + results = await _save_round(search_service, file_service, entity, round_start) + failures += [r for r in results if isinstance(r, BaseException)] + if failures: + break + + messages = " ".join(str(failure) for failure in failures) + assert failures, "without the lock, overlapping saves were expected to collide" + assert "search_index_pkey" in messages or "deadlock detected" in messages + + +async def _refresh_relations_concurrently(search_service, entity: Entity) -> list[BaseException]: + """Refresh with the content in hand, as the batch indexer does; relation rows ride along. + + The round starts from an entity with no rows, as after a delete another writer has + committed. Every refresh's delete then finds nothing to wait on, and all of them + reach their inserts together: the shape that produced ``Key (id, relation, N)``. + """ + await search_service.repository.delete_by_entity_id(entity.id) + results = await asyncio.gather( + *( + search_service.index_entity_data(entity, content=_long_body(v)) + for v in range(CONCURRENT_SAVES) + ), + return_exceptions=True, + ) + return [r for r in results if isinstance(r, BaseException)] + + +async def test_concurrent_refreshes_keep_relation_rows_single( + monkeypatch, db_backend, search_service, full_entity, entity_repository, session_maker +): + _skip_unless_postgres(db_backend) + entity = await _reload(session_maker, entity_repository, full_entity) + await search_service.index_entity_data(entity, content="before") + before = await _projection(search_service, entity) + assert [kind for kind, _ in before].count("relation") == 2 + _widen_the_window(monkeypatch, search_service.repository) + + for _ in range(3): + assert await _refresh_relations_concurrently(search_service, entity) == [] + assert await _projection(search_service, entity) == before + + +async def test_negative_control_relation_rows_collide_without_the_lock( + monkeypatch, db_backend, search_service, full_entity, entity_repository, session_maker +): + _skip_unless_postgres(db_backend) + entity = await _reload(session_maker, entity_repository, full_entity) + await search_service.index_entity_data(entity, content="before") + _widen_the_window(monkeypatch, search_service.repository) + _remove_the_lock(monkeypatch) + + failures: list[BaseException] = [] + for _ in range(10): + failures += await _refresh_relations_concurrently(search_service, entity) + if failures: + break + assert failures, "without the lock, overlapping refreshes were expected to collide" + assert "search_index_pkey" in " ".join(str(failure) for failure in failures) + + +async def test_an_accepted_write_and_an_index_refresh_take_turns( + monkeypatch, db_backend, search_service, full_entity, entity_repository, session_maker +): + """The save path's accepted write rewrites the entity row too; it waits its turn.""" + _skip_unless_postgres(db_backend) + entity = await _reload(session_maker, entity_repository, full_entity) + await search_service.index_entity_data(entity, content="before") + _widen_the_window(monkeypatch, search_service.repository) + accepted = AcceptedNoteSearchRepository(project_id=entity.project_id) + + async def accepted_write(version: int) -> None: + body = _long_body(version) + async with db.scoped_session(session_maker) as session: + await accepted.refresh_entity( + session, + AcceptedNoteSearchRow( + id=entity.id, + title=entity.title, + content_stems=body, + content_snippet=body, + permalink=entity.permalink, + file_path=entity.file_path, + item_type="entity", + note_type=entity.note_type, + entity_id=entity.id, + created_at=entity.created_at, + updated_at=entity.updated_at, + project_id=entity.project_id, + ), + ) + await asyncio.sleep(0.05) + + for round_start in (0, 100, 200): + results = await asyncio.gather( + *( + accepted_write(v) + if v % 2 + else search_service.index_entity_data(entity, content=_long_body(v)) + for v in range(round_start, round_start + CONCURRENT_SAVES) + ), + return_exceptions=True, + ) + assert [r for r in results if isinstance(r, BaseException)] == [] + rows = await search_service.repository.get_entity_search_rows(entity.id) + assert [row.type for row in rows].count("entity") == 1