diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index f49effeb..d08b31af 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -395,7 +395,8 @@ async def _refresh_doc_metadata( updated = True if updated: - result = await client.document_repository.update_meta(doc) + async with client.store._write_lock: + result = await client.document_repository.update_meta(doc) # Reclaim the document_meta churn from rolling source_revision sweeps. # The vacuum is debounced, and document_meta is tiny, so this is cheap. if client._config.storage.auto_vacuum: diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index e0986add..35b9cffe 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -520,7 +520,7 @@ class Store: if self._vacuum_lock.locked(): return - async with self._vacuum_lock: + async with self._vacuum_lock, self._write_lock: try: # Evaluate config at runtime to allow dynamic changes if retention_seconds is None: @@ -866,6 +866,9 @@ class Store: async def create_tag(self, name: str) -> None: """Tag the current version of every table with the given name. + Serializes with client writes via the write lock so a write cannot + land between the version snapshot and the per-table tag creation. + Raises: ReadOnlyError: If the store is in read-only mode. ValueError: If the tag already exists on any table. A partial tag @@ -875,30 +878,31 @@ class Store: self._assert_writable() tables = self._tables() - existing = [ - table_name - for table_name, table in tables.items() - if name in await table.tags.list() - ] - if len(existing) == len(tables): - raise ValueError(f"Tag '{name}' already exists") - if existing: - raise ValueError( - f"Tag '{name}' already exists on some tables " - f"({', '.join(existing)}); delete it first with delete_tag" - ) + async with self._write_lock: + existing = [ + table_name + for table_name, table in tables.items() + if name in await table.tags.list() + ] + if len(existing) == len(tables): + raise ValueError(f"Tag '{name}' already exists") + if existing: + raise ValueError( + f"Tag '{name}' already exists on some tables " + f"({', '.join(existing)}); delete it first with delete_tag" + ) - versions = await self.current_table_versions() - created: list[str] = [] - try: - for table_name, table in tables.items(): - await table.tags.create(name, versions[table_name]) - created.append(table_name) - except Exception: - for table_name in created: - with contextlib.suppress(Exception): - await tables[table_name].tags.delete(name) - raise + versions = await self.current_table_versions() + created: list[str] = [] + try: + for table_name, table in tables.items(): + await table.tags.create(name, versions[table_name]) + created.append(table_name) + except Exception: + for table_name in created: + with contextlib.suppress(Exception): + await tables[table_name].tags.delete(name) + raise async def list_tags(self) -> dict[str, TagInfo]: """Aggregate per-table tags into database-level tags. @@ -920,18 +924,21 @@ class Store: async def delete_tag(self, name: str) -> None: """Delete the tag from every table that has it. + Serializes with create_tag and client writes via the write lock. + Raises: ReadOnlyError: If the store is in read-only mode. ValueError: If no table has the tag. """ self._assert_writable() - found = False - for table in self._tables().values(): - if name in await table.tags.list(): - await table.tags.delete(name) - found = True - if not found: - raise ValueError(f"Tag '{name}' does not exist") + async with self._write_lock: + found = False + for table in self._tables().values(): + if name in await table.tags.list(): + await table.tags.delete(name) + found = True + if not found: + raise ValueError(f"Tag '{name}' does not exist") async def _checkout_tables_before(self, before: datetime) -> None: """Checkout all tables to their state at or before the given datetime. diff --git a/tests/store/test_tags.py b/tests/store/test_tags.py index e8d010e4..70eaf03a 100644 --- a/tests/store/test_tags.py +++ b/tests/store/test_tags.py @@ -85,6 +85,37 @@ async def test_create_tag_rolls_back_own_tags_on_failure(temp_db_path, monkeypat assert tags["keep"].complete is True +@pytest.mark.asyncio +async def test_create_tag_waits_for_write_lock(temp_db_path): + """create_tag serializes with client writes so a write cannot land + between the version snapshot and the per-table tag creation.""" + async with Store(temp_db_path, create=True) as store: + async with store._write_lock: + task = asyncio.create_task(store.create_tag("release-1")) + await asyncio.sleep(0.1) + assert not task.done() + await task + + tags = await store.list_tags() + assert tags["release-1"].complete is True + + +@pytest.mark.asyncio +async def test_delete_tag_waits_for_write_lock(temp_db_path): + """delete_tag serializes with create_tag and client writes so it cannot + remove tags out from under a concurrent create_tag.""" + async with Store(temp_db_path, create=True) as store: + await store.create_tag("release-1") + + async with store._write_lock: + task = asyncio.create_task(store.delete_tag("release-1")) + await asyncio.sleep(0.1) + assert not task.done() + await task + + assert await store.list_tags() == {} + + @pytest.mark.asyncio async def test_delete_tag(temp_db_path): async with Store(temp_db_path, create=True) as store: @@ -144,6 +175,18 @@ async def test_vacuum_cleans_untagged_versions_and_keeps_tagged(temp_db_path): assert rows == 2 +@pytest.mark.asyncio +async def test_vacuum_waits_for_write_lock(temp_db_path): + """Vacuum serializes with writers and tag operations so a tag cannot be + created between _tag_safe_retention's read and the optimize call.""" + async with Store(temp_db_path, create=True) as store: + async with store._write_lock: + task = asyncio.create_task(store.vacuum(retention_seconds=0)) + await asyncio.sleep(0.1) + assert not task.done() + await task + + @pytest.mark.asyncio async def test_tag_writes_raise_when_read_only(temp_db_path): async with Store(temp_db_path, create=True) as store: diff --git a/tests/test_vacuum_debounce.py b/tests/test_vacuum_debounce.py index 2b4176f8..83a96dbc 100644 --- a/tests/test_vacuum_debounce.py +++ b/tests/test_vacuum_debounce.py @@ -92,3 +92,33 @@ async def test_metadata_refresh_sweep_schedules_vacuum(temp_db_path): source_metadata={"source_revision": "r2", "md5": "same"}, ) assert client._vacuum_dirty is True + + +@pytest.mark.asyncio +async def test_metadata_refresh_waits_for_write_lock(temp_db_path): + """The revision/MD5 short-circuit write serializes with other writers so + it cannot land inside another writer's critical section (e.g. between + create_tag's version snapshot and its per-table tag creation).""" + dim = Config.embeddings.model.vector_dim + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.import_document( + _docling_doc("d", "body"), + [Chunk(content="body", embedding=[0.1] * dim, order=0)], + uri="mem://sweep", + metadata={"source_revision": "r1"}, + ) + + async with client.store._write_lock: + task = asyncio.create_task( + _refresh_doc_metadata( + client, + doc, + title=None, + user_metadata={}, + source_metadata={"source_revision": "r2", "md5": "same"}, + ) + ) + await asyncio.sleep(0.1) + assert not task.done() + refreshed = await task + assert refreshed.metadata["source_revision"] == "r2"