Serialize tag operations, vacuum, and metadata refresh with writes

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:

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 = [
+        async with self._write_lock:
-            table_name
+            existing = [
-            for table_name, table in tables.items()
+                table_name
-            if name in await table.tags.list()
+                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 len(existing) == len(tables):
-        if existing:
+                raise ValueError(f"Tag '{name}' already exists")
-            raise ValueError(
+            if existing:
-                f"Tag '{name}' already exists on some tables "
+                raise ValueError(
-                f"({', '.join(existing)}); delete it first with delete_tag"
+                    f"Tag '{name}' already exists on some tables "
-            )
+                    f"({', '.join(existing)}); delete it first with delete_tag"
+                )
 
-        versions = await self.current_table_versions()
+            versions = await self.current_table_versions()
-        created: list[str] = []
+            created: list[str] = []
-        try:
+            try:
-            for table_name, table in tables.items():
+                for table_name, table in tables.items():
-                await table.tags.create(name, versions[table_name])
+                    await table.tags.create(name, versions[table_name])
-                created.append(table_name)
+                    created.append(table_name)
-        except Exception:
+            except Exception:
-            for table_name in created:
+                for table_name in created:
-                with contextlib.suppress(Exception):
+                    with contextlib.suppress(Exception):
-                    await tables[table_name].tags.delete(name)
+                        await tables[table_name].tags.delete(name)
-            raise
+                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
+        async with self._write_lock:
-        for table in self._tables().values():
+            found = False
-            if name in await table.tags.list():
+            for table in self._tables().values():
-                await table.tags.delete(name)
+                if name in await table.tags.list():
-                found = True
+                    await table.tags.delete(name)
-        if not found:
+                    found = True
-            raise ValueError(f"Tag '{name}' does not exist")
+            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.

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:

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"