improvement(knowledge): bound each reconciliation window once and read it through the v2 id range

Author: waleedlatif1

apps/sim/lib/knowledge/__integration__/listing-continuation.integration.ts

@@ -23,7 +23,7 @@ import { toNumberOrNull, toStringOrNull } from '@sim/utils/coerce'
 import { generateId } from '@sim/utils/id'
 import { toArray, toRecord } from '@sim/utils/object'
 import { and, eq, inArray, sql } from 'drizzle-orm'
-import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const fixture = vi.hoisted(() => ({
   storageRoot: '',
@@ -901,11 +901,22 @@ describe('durable source and member cycles in PostgreSQL', () => {
     }
   })
 
-  it('bounds every reconciliation page by the ids it scans when absence is rare and late in id order', async () => {
-    const connectorId = generateId()
-    const runId = generateId()
-    const startedAt = new Date()
+  describe('reconciliation windows', () => {
+    const UNRELATED_FILENAME = 'reconciliation-window-unrelated'
     const acl = [`u:${ids.aliceId}@fixture.test`]
+    let connectorId = ''
+    let runId = ''
+    let startedAt = new Date()
+    beforeEach(() => {
+      connectorId = generateId()
+      runId = generateId()
+      startedAt = new Date()
+    })
+    afterEach(async () => {
+      await db.delete(document).where(eq(document.filename, UNRELATED_FILENAME))
+      await db.delete(document).where(eq(document.connectorId, connectorId))
+      await db.delete(knowledgeConnector).where(eq(knowledgeConnector.id, connectorId))
+    })
     const row = (id: string) => ({
       id,
       knowledgeBaseId: ids.knowledgeBaseId,
@@ -920,79 +931,142 @@ describe('durable source and member cycles in PostgreSQL', () => {
       acl,
       aclVerifiedAt: startedAt,
     })
-    /** Ids sort present rows first, so each absent row is found only after three full windows. */
-    const present = Array.from({ length: 3 * RECONCILIATION_WINDOW_SIZE }, (_, index) => ({
-      ...row(`${connectorId}-a-${String(index).padStart(6, '0')}`),
-      sourceSeenAt: startedAt,
-    }))
-    const absent = Array.from({ length: 3 }, (_, index) => ({
-      ...row(`${connectorId}-z-live-${index}`),
-      sourceSeenAt: null,
-    }))
-    const tombstoned = Array.from({ length: 2 }, (_, index) => ({
-      ...row(`${connectorId}-z-tombstone-${index}`),
-      sourceSeenAt: null,
-      deletedAt: new Date(startedAt.getTime() - 60_000),
-    }))
-    const checkpoint = beginListingCheckpoint({
-      fingerprint: listingFingerprint({ connectorId }),
-      generationId: runId,
-      startedAt,
-    })
-    checkpoint.complete = true
-    checkpoint.listedCount = present.length
-    await db.insert(knowledgeConnector).values({
-      id: connectorId,
-      knowledgeBaseId: ids.knowledgeBaseId,
-      connectorType: 'google_drive',
-      sourceConfig: {},
-      accessMode: 'admin',
-      status: 'syncing',
-      syncLockToken: runId,
-      listingCheckpoint: checkpoint,
-    })
-    const rows = [...present, ...absent, ...tombstoned]
-    for (let offset = 0; offset < rows.length; offset += 1_000)
-      await db.insert(document).values(rows.slice(offset, offset + 1_000))
-    await db.execute(sql`ANALYZE document`)
-    const hardDelete = vi.spyOn(documentService, 'hardDeleteDocuments')
-    const statements = vi.spyOn(db.$client, 'unsafe')
-    try {
-      const [connector] = await db
-        .select()
-        .from(knowledgeConnector)
-        .where(eq(knowledgeConnector.id, connectorId))
-      const stats = result()
-      const pass = await runConnectorContentPass({
-        connectorId,
-        connector,
-        connectorConfig: CONNECTOR_REGISTRY.google_drive,
-        sourceConfig: {},
-        syncContext: {},
-        kbOwner: { userId: ids.aliceId, workspaceId: ids.workspaceId },
-        billingAttribution: billing,
-        result: stats,
-        lease: createContentSyncLease(connectorId, runId),
-        leaseKind: 'content',
-        runId,
+    /**
+     * Seeds a connector whose listing already completed, so the pass goes straight to
+     * reconciliation. Other documents spread across the id space make the connector the minority
+     * it is at scale, where the planner reads a window through the connector's own id range
+     * rather than a primary-key range over every connector's rows.
+     */
+    const seed = async (rows: ReturnType<typeof row>[], listedCount: number) => {
+      await db.execute(sql`
+        INSERT INTO ${document} (id, knowledge_base_id, filename, file_url, file_size, mime_type, processing_status)
+        SELECT md5(${connectorId} || g), ${ids.knowledgeBaseId}, ${UNRELATED_FILENAME}, '', 0, 'text/plain', 'completed'
+        FROM generate_series(1, 20 * ${RECONCILIATION_WINDOW_SIZE}) g`)
+      const checkpoint = beginListingCheckpoint({
         fingerprint: listingFingerprint({ connectorId }),
-        documentAccess: 'admin',
-        getAccessToken: async () => 'fixture',
-        hydration: { getDocument: fixture.get },
-        forceRehydrate: false,
-        deadlineAt: Date.now() + 30_000,
+        generationId: runId,
+        startedAt,
       })
-      const walked = statements.mock.calls
-        .map(([query, params]) => ({ query, params }))
-        .filter(({ query }) =>
-          /^select .* from "document" .*order by "document"\."id"/is.test(query)
-        )
-      statements.mockRestore()
+      checkpoint.complete = true
+      checkpoint.listedCount = listedCount
+      await db.insert(knowledgeConnector).values({
+        id: connectorId,
+        knowledgeBaseId: ids.knowledgeBaseId,
+        connectorType: 'google_drive',
+        sourceConfig: {},
+        accessMode: 'admin',
+        status: 'syncing',
+        syncLockToken: runId,
+        listingCheckpoint: checkpoint,
+      })
+      for (let offset = 0; offset < rows.length; offset += 1_000)
+        await db.insert(document).values(rows.slice(offset, offset + 1_000))
+      await db.execute(sql`ANALYZE document`)
+    }
+    const isWindowBound = (query: string) =>
+      /^select .* from "document" .*order by "document"\."id" .*offset \$\d+$/is.test(query)
+    const isWindowPage = (query: string) => /^\s*with "document" as materialized/i.test(query)
+    /** Runs the pass, returning the window statements it issued outside transactions. */
+    const reconcile = async () => {
+      const hardDelete = vi.spyOn(documentService, 'hardDeleteDocuments')
+      const statements = vi.spyOn(db.$client, 'unsafe')
+      try {
+        const [connector] = await db
+          .select()
+          .from(knowledgeConnector)
+          .where(eq(knowledgeConnector.id, connectorId))
+        const stats = result()
+        const pass = await runConnectorContentPass({
+          connectorId,
+          connector,
+          connectorConfig: CONNECTOR_REGISTRY.google_drive,
+          sourceConfig: {},
+          syncContext: {},
+          kbOwner: { userId: ids.aliceId, workspaceId: ids.workspaceId },
+          billingAttribution: billing,
+          result: stats,
+          lease: createContentSyncLease(connectorId, runId),
+          leaseKind: 'content',
+          runId,
+          fingerprint: listingFingerprint({ connectorId }),
+          documentAccess: 'admin',
+          getAccessToken: async () => 'fixture',
+          hydration: { getDocument: fixture.get },
+          forceRehydrate: false,
+          deadlineAt: Date.now() + 60_000,
+        })
+        return {
+          pass,
+          stats,
+          hardDeleted: hardDelete.mock.calls.flatMap(([batch]) => batch),
+          walked: statements.mock.calls
+            .map(([query, params]) => ({ query, params }))
+            .filter(({ query }) => isWindowBound(query) || isWindowPage(query)),
+        }
+      } finally {
+        statements.mockRestore()
+        hardDelete.mockRestore()
+      }
+    }
+    /**
+     * Plans a walk statement as a large table would, on an index rather than a sequential scan,
+     * returning the document rows it read and the indexes it read them through.
+     */
+    const explainWalk = async (query: string, params: Parameters<typeof db.$client.unsafe>[1]) => {
+      const [explained] = await db.$client.begin(async (tx) => {
+        await tx.unsafe('SET LOCAL enable_seqscan = off')
+        return tx.unsafe(`EXPLAIN (ANALYZE, FORMAT JSON) ${query}`, params)
+      })
+      const nodes = (node: unknown): Record<string, unknown>[] => {
+        const plan = toRecord(node)
+        return [plan, ...toArray(plan.Plans).flatMap(nodes)]
+      }
+      const plan = nodes(toRecord(toArray(explained['QUERY PLAN'])[0]).Plan)
+      const scans = plan.filter((node) => node['Relation Name'] === 'document')
+      return {
+        scanned: scans.reduce(
+          (total, plan) =>
+            total +
+            ((toNumberOrNull(plan['Actual Rows']) ?? 0) +
+              (toNumberOrNull(plan['Rows Removed by Filter']) ?? 0)) *
+              (toNumberOrNull(plan['Actual Loops']) ?? 1),
+          0
+        ),
+        indexes: plan.map((node) => toStringOrNull(node['Index Name'])).filter(Boolean),
+      }
+    }
+    /** Every window statement reads through the connector's v2 id range, at most one window of it. */
+    const expectBounded = async (
+      walked: { query: string; params: Parameters<typeof db.$client.unsafe>[1] }[]
+    ) => {
+      expect(walked.length).toBeGreaterThan(0)
+      for (const { query, params } of walked) {
+        const plan = await explainWalk(query, params)
+        expect(plan.indexes, query).toEqual(['doc_connector_reconciliation_v2_idx'])
+        expect(plan.scanned, query).toBeLessThanOrEqual(RECONCILIATION_WINDOW_SIZE)
+      }
+    }
+
+    it('bounds every page by the ids it scans when absence is rare and late in id order', async () => {
+      /** Ids sort present rows first, so each absent row is found only after three full windows. */
+      const present = Array.from({ length: 3 * RECONCILIATION_WINDOW_SIZE }, (_, index) => ({
+        ...row(`${connectorId}-a-${String(index).padStart(6, '0')}`),
+        sourceSeenAt: startedAt,
+      }))
+      const absent = Array.from({ length: 3 }, (_, index) => ({
+        ...row(`${connectorId}-z-live-${index}`),
+        sourceSeenAt: null,
+      }))
+      const tombstoned = Array.from({ length: 2 }, (_, index) => ({
+        ...row(`${connectorId}-z-tombstone-${index}`),
+        sourceSeenAt: null,
+        deletedAt: new Date(startedAt.getTime() - 60_000),
+      }))
+      await seed([...present, ...absent, ...tombstoned], present.length)
+      const { pass, stats, hardDeleted, walked } = await reconcile()
       expect(pass).toMatchObject({ complete: true, holdNotice: null })
       expect(stats.docsDeleted).toBe(absent.length)
-      expect(hardDelete.mock.calls.flatMap(([batch]) => batch).sort()).toEqual(
-        tombstoned.map((item) => item.id).sort()
-      )
+      expect(hardDeleted.sort()).toEqual(tombstoned.map((item) => item.id).sort())
       const stored = await db
         .select({ id: document.id, acl: document.acl, deletedAt: document.deletedAt })
         .from(document)
@@ -1002,45 +1076,44 @@ describe('durable source and member cycles in PostgreSQL', () => {
       expect(
         stored.filter((item) => item.id.includes('-z-')).every((item) => item.acl.length === 0)
       ).toBe(true)
+      await expectBounded(walked)
+    }, 120_000)
+
+    it('bounds a dense window once and pages it on the id range, not the tombstone index', async () => {
       /**
-       * Plans every walk statement as a large table would, on an index rather than a sequential
-       * scan. This connector dominates the fixture table, so the planner may walk the primary key
-       * and filter out the other connectors' rows; only those may exceed a window.
+       * Three hard-delete pages of absent tombstones open the first window; the rest of the
+       * connector is tombstones the listing still sees, which the hard walk must pass over.
        */
-      const [{ foreign }] = await db
-        .select({ foreign: sql<number>`count(*)::int` })
-        .from(document)
-        .where(sql`${document.connectorId} IS DISTINCT FROM ${connectorId}`)
-      expect(walked.length).toBeGreaterThan(0)
-      for (const { query, params } of walked) {
-        const [explained] = await db.$client.begin(async (tx) => {
-          await tx.unsafe('SET LOCAL enable_seqscan = off')
-          return tx.unsafe(`EXPLAIN (ANALYZE, FORMAT JSON) ${query}`, params)
-        })
-        const scanned = (node: unknown): number => {
-          const plan = toRecord(node)
-          const own =
-            toStringOrNull(plan['Relation Name']) === 'document'
-              ? ((toNumberOrNull(plan['Actual Rows']) ?? 0) +
-                  (toNumberOrNull(plan['Rows Removed by Filter']) ?? 0)) *
-                (toNumberOrNull(plan['Actual Loops']) ?? 1)
-              : 0
-          return (
-            own + toArray(plan.Plans).reduce<number>((total, child) => total + scanned(child), 0)
-          )
-        }
-        expect(
-          scanned(toRecord(toArray(explained['QUERY PLAN'])[0]).Plan),
-          query
-        ).toBeLessThanOrEqual(RECONCILIATION_WINDOW_SIZE + foreign)
-      }
-    } finally {
-      statements.mockRestore()
-      hardDelete.mockRestore()
-      await db.delete(document).where(eq(document.connectorId, connectorId))
-      await db.delete(knowledgeConnector).where(eq(knowledgeConnector.id, connectorId))
-    }
-  }, 120_000)
+      const dense = Array.from({ length: 75 }, (_, index) => ({
+        ...row(`${connectorId}-a-${String(index).padStart(3, '0')}`),
+        acl: [],
+        sourceSeenAt: null,
+        deletedAt: new Date(startedAt.getTime() - 60_000),
+      }))
+      const seenTombstones = Array.from({ length: 3 * RECONCILIATION_WINDOW_SIZE }, (_, index) => ({
+        ...row(`${connectorId}-b-${String(index).padStart(6, '0')}`),
+        sourceSeenAt: startedAt,
+        deletedAt: new Date(startedAt.getTime() - 60_000),
+      }))
+      await seed([...dense, ...seenTombstones], seenTombstones.length)
+      const { pass, hardDeleted, walked } = await reconcile()
+      expect(pass).toMatchObject({ complete: true, holdNotice: null })
+      expect(hardDeleted.sort()).toEqual(dense.map((item) => item.id).sort())
+      expect(
+        await db
+          .select({ id: document.id })
+          .from(document)
+          .where(eq(document.connectorId, connectorId))
+      ).toHaveLength(seenTombstones.length)
+      /** The only walk is the hard one: one bound per window, the dense first one included, then the tail. */
+      expect(walked.filter(({ query }) => isWindowBound(query))).toHaveLength(
+        Math.floor((dense.length + seenTombstones.length) / RECONCILIATION_WINDOW_SIZE) + 1
+      )
+      expect(walked.filter(({ query }) => isWindowPage(query)).length).toBeGreaterThan(3)
+      /** `deleted_at < $1` implies the tombstone index, which would read every connector tombstone. */
+      await expectBounded(walked)
+    }, 120_000)
+  })
 
   it('indexes a page, resumes under a new lease, and reconciles absence only after EOF', async () => {
     await db

apps/sim/lib/knowledge/__integration__/scale.integration.ts

@@ -337,16 +337,12 @@ describe.skipIf(!enabled)('knowledge scale: isolated real PostgreSQL, no provide
       'doc_connector_reconciliation_v2_idx'
     )
     const [end] = await windowEnd
-    const windowPage = db
-      .select({ id: document.id })
-      .from(document)
-      .where(
-        sql`${owned} AND (source_seen_at IS NULL OR source_seen_at < ${startedAt.toISOString()}::timestamp)
-        AND cardinality(acl) > 0 AND id <= ${end.id}`
-      )
-      .orderBy(asc(document.id))
-      .limit(PAGE_SIZE)
-    await explain('reconciliation.window.page', windowPage.getSQL())
+    const windowPage = sql`WITH document AS MATERIALIZED (
+        SELECT id, source_seen_at, acl FROM document WHERE ${owned} AND id <= ${end.id})
+      SELECT id FROM document
+      WHERE (source_seen_at IS NULL OR source_seen_at < ${startedAt.toISOString()}::timestamp) AND cardinality(acl) > 0
+      ORDER BY id LIMIT ${PAGE_SIZE}`
+    await explain('reconciliation.window.page', windowPage)
     expect(planIndexNames('reconciliation.window.page')).toContain(
       'doc_connector_reconciliation_v2_idx'
     )

apps/sim/lib/knowledge/connectors/member-observations.test.ts

@@ -339,8 +339,7 @@ describe('applyMemberDocumentLifecycle', () => {
   it('reports a reclaimed lease during a purge batch as the run being superseded', async () => {
     dbChainMockFns.returning.mockResolvedValueOnce([])
     queueTableRows(schemaMock.document, [])
-    /** The resurrection walk's window bound and its empty page. */
-    queueTableRows(schemaMock.document, [])
+    /** The resurrection walk's window bound; its page, read through `db.execute`, is empty. */
     queueTableRows(schemaMock.document, [])
     queueTableRows(schemaMock.document, [{ id: 'd-1' }])
     vi.mocked(hardDeleteDocuments).mockRejectedValueOnce(