Merge pull request #393 from mcdonc/perf/batch-sync-state-writes

perf: batch sync_state writes during poller sweeps

haiku_rag_slim/haiku/rag/ingester/pollers/base.py

@@ -5,7 +5,7 @@ from datetime import UTC, datetime
 
 from haiku.rag.config import SourceConfig
 from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
-from haiku.rag.ingester.queue.models import JobOp
+from haiku.rag.ingester.queue.models import JobOp, SyncRow
 from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
 from haiku.rag.ingester.sources.base import (
     Source,
@@ -129,11 +129,13 @@ class BasePoller:
                     SourceEventKind.DELETE: 0,
                     SourceEventKind.UNCHANGED: 0,
                 }
+                sync_batch: list[SyncRow] = []
                 async for event in self.source.discover(
                     since=revisions, known_uris=known
                 ):
                     counts[event.kind] += 1
-                    await self._handle_event(event)
+                    await self._handle_event(event, sync_batch)
+                await self._sync.batch_upsert(sync_batch)
                 self._breaker.record_success()
                 self._last_polled_at = datetime.now(UTC)
                 self._last_skip_reason = None
@@ -163,7 +165,11 @@ class BasePoller:
                 )
                 return False
 
-    async def _handle_event(self, event: SourceEvent) -> None:
+    async def _handle_event(
+        self,
+        event: SourceEvent,
+        sync_batch: list[SyncRow],
+    ) -> None:
         if event.kind is SourceEventKind.UPSERT:
             await self._jobs.enqueue(
                 event.source_id,
@@ -176,19 +182,11 @@ class BasePoller:
             # Don't write revision to sync_state here — the worker writes it
             # after a successful ingestion. last_seen_at gets bumped to keep
             # orphan detection accurate.
-            await self._sync.upsert(
-                event.source_id,
-                event.uri,
-                revision=None,
-                content_hash=None,
-            )
+            sync_batch.append(SyncRow(event.source_id, event.uri, None, None, False))
         elif event.kind is SourceEventKind.UNCHANGED:
             # Touch last_seen_at without changing the stored revision.
-            await self._sync.upsert(
-                event.source_id,
-                event.uri,
-                revision=event.revision,
-                content_hash=None,
+            sync_batch.append(
+                SyncRow(event.source_id, event.uri, event.revision, None, False)
             )
         elif event.kind is SourceEventKind.DELETE:
             if not self.config.delete_orphans:

haiku_rag_slim/haiku/rag/ingester/queue/models.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 from enum import StrEnum
+from typing import NamedTuple
 
 from pydantic import BaseModel
 
@@ -44,3 +45,14 @@ class SyncStateRow(BaseModel):
     content_hash: str | None = None
     last_seen_at: datetime
     last_ingested_at: datetime | None = None
+
+
+class SyncRow(NamedTuple):
+    """A pending sync_state write collected during a sweep and flushed in one
+    transaction by SyncStateRepo.batch_upsert()."""
+
+    source_id: str
+    uri: str
+    revision: str | None
+    content_hash: str | None
+    ingested: bool

haiku_rag_slim/haiku/rag/ingester/queue/repository.py

@@ -5,7 +5,13 @@ from datetime import UTC, datetime, timedelta
 
 import aiosqlite
 
-from haiku.rag.ingester.queue.models import Job, JobOp, JobStatus, SyncStateRow
+from haiku.rag.ingester.queue.models import (
+    Job,
+    JobOp,
+    JobStatus,
+    SyncRow,
+    SyncStateRow,
+)
 
 
 def _utcnow_iso() -> str:
@@ -504,6 +510,33 @@ class SyncStateRepo:
                 pass
             await self._conn.commit()
 
+    async def batch_upsert(self, rows: list[SyncRow]) -> None:
+        """Batch insert-or-update sync_state rows in a single transaction."""
+        if not rows:
+            return
+        now = _utcnow_iso()
+        async with self._lock:
+            for source_id, uri, revision, content_hash, ingested in rows:
+                ingested_at = now if ingested else None
+                await self._conn.execute(
+                    """
+                    INSERT INTO sync_state (
+                        source_id, uri, revision, content_hash,
+                        last_seen_at, last_ingested_at
+                    )
+                    VALUES (?, ?, ?, ?, ?, ?)
+                    ON CONFLICT(source_id, uri) DO UPDATE SET
+                        revision = COALESCE(excluded.revision, revision),
+                        content_hash = COALESCE(excluded.content_hash, content_hash),
+                        last_seen_at = excluded.last_seen_at,
+                        last_ingested_at = COALESCE(
+                            excluded.last_ingested_at, last_ingested_at
+                        )
+                    """,
+                    (source_id, uri, revision, content_hash, now, ingested_at),
+                )
+            await self._conn.commit()
+
     async def delete(self, source_id: str, uri: str) -> None:
         async with self._lock:
             async with self._conn.execute(

tests/ingester/test_queue.py

@@ -5,7 +5,7 @@ import aiosqlite
 import pytest
 
 from haiku.rag.ingester.queue.migrations import apply_migrations, open_queue
-from haiku.rag.ingester.queue.models import JobOp, JobStatus
+from haiku.rag.ingester.queue.models import JobOp, JobStatus, SyncRow
 from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
 
 
@@ -834,3 +834,60 @@ async def test_sync_state_upsert_replaces_revision_when_provided(sync):
     assert row.revision == "v2"
     assert row.content_hash == "hash-v2"
     assert row.last_ingested_at is not None
+
+
+@pytest.mark.asyncio
+async def test_sync_state_batch_upsert_inserts_multiple_rows(sync):
+    """batch_upsert writes many rows in a single transaction."""
+    await sync.batch_upsert(
+        [
+            SyncRow("s", "u1", "rev1", "hash1", False),
+            SyncRow("s", "u2", "rev2", "hash2", False),
+            SyncRow("s", "u3", "rev3", None, True),
+        ]
+    )
+    assert await sync.get_revision_snapshot("s") == {
+        "u1": "rev1",
+        "u2": "rev2",
+        "u3": "rev3",
+    }
+    row = await sync.get_row("s", "u3")
+    assert row is not None
+    assert row.last_ingested_at is not None
+
+
+@pytest.mark.asyncio
+async def test_sync_state_batch_upsert_updates_existing(sync):
+    """batch_upsert applies ON CONFLICT update semantics like upsert()."""
+    await sync.upsert("s", "u1", revision="old", content_hash="old-hash")
+    await sync.batch_upsert(
+        [
+            SyncRow("s", "u1", "new", "new-hash", False),
+        ]
+    )
+    row = await sync.get_row("s", "u1")
+    assert row is not None
+    assert row.revision == "new"
+    assert row.content_hash == "new-hash"
+
+
+@pytest.mark.asyncio
+async def test_sync_state_batch_upsert_preserves_revision_when_none(sync):
+    """batch_upsert with revision=None leaves existing revision in place."""
+    await sync.upsert("s", "u1", revision="keep", content_hash="keep-hash")
+    await sync.batch_upsert(
+        [
+            SyncRow("s", "u1", None, None, False),
+        ]
+    )
+    row = await sync.get_row("s", "u1")
+    assert row is not None
+    assert row.revision == "keep"
+    assert row.content_hash == "keep-hash"
+
+
+@pytest.mark.asyncio
+async def test_sync_state_batch_upsert_empty_is_noop(sync):
+    """batch_upsert with an empty list does nothing."""
+    await sync.batch_upsert([])
+    assert await sync.get_revision_snapshot("s") == {}