perf(knowledge): reconcile a completed listing in one pass over its windows

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

@@ -966,8 +966,8 @@ describe('durable source and member cycles in PostgreSQL', () => {
       await db.execute(sql`ANALYZE document`)
     }
     const isWindowStatement = (query: string) => /^\s*with "document" as materialized/i.test(query)
-    const isWindowCount = (query: string) => query.includes('json_build_array')
-    /** Runs the pass, returning the window scans and counts it issued outside transactions. */
+    const isWindowTally = (query: string) => query.includes('json_build_array')
+    /** Runs the pass, returning the window walks and tallies it issued outside transactions. */
     const reconcile = async () => {
       const hardDelete = vi.spyOn(documentService, 'hardDeleteDocuments')
       const statements = vi.spyOn(db.$client, 'unsafe')
@@ -1002,10 +1002,10 @@ describe('durable source and member cycles in PostgreSQL', () => {
           hardDeleted: hardDelete.mock.calls.flatMap(([batch]) => batch),
           walked: statements.mock.calls
             .map(([query, params]) => ({ query, params }))
-            .filter(({ query }) => isWindowStatement(query) && !isWindowCount(query)),
-          counted: statements.mock.calls
+            .filter(({ query }) => isWindowStatement(query) && !isWindowTally(query)),
+          tallied: statements.mock.calls
             .map(([query, params]) => ({ query, params }))
-            .filter(({ query }) => isWindowStatement(query) && isWindowCount(query)),
+            .filter(({ query }) => isWindowStatement(query) && isWindowTally(query)),
         }
       } finally {
         statements.mockRestore()
@@ -1042,10 +1042,10 @@ describe('durable source and member cycles in PostgreSQL', () => {
      * must never happen.
      */
     const expectBounded = async (
-      walked: { query: string; params: Parameters<typeof db.$client.unsafe>[1] }[]
+      statements: { query: string; params: Parameters<typeof db.$client.unsafe>[1] }[]
     ) => {
-      expect(walked.length).toBeGreaterThan(0)
-      for (const { query, params } of walked) {
+      expect(statements.length).toBeGreaterThan(0)
+      for (const { query, params } of statements) {
         const { read, scans } = await explainWalk(query, params)
         expect(read, `${query}\n${JSON.stringify(scans, null, 2)}`).toBeLessThanOrEqual(
           RECONCILIATION_WINDOW_SIZE
@@ -1069,7 +1069,7 @@ describe('durable source and member cycles in PostgreSQL', () => {
         deletedAt: new Date(startedAt.getTime() - 60_000),
       }))
       await seed([...present, ...absent, ...tombstoned], present.length)
-      const { pass, stats, hardDeleted, walked, counted } = await reconcile()
+      const { pass, stats, hardDeleted, walked, tallied } = await reconcile()
       expect(pass).toMatchObject({ complete: true, holdNotice: null })
       expect(stats.docsDeleted).toBe(absent.length)
       expect(hardDeleted.sort()).toEqual(tombstoned.map((item) => item.id).sort())
@@ -1082,14 +1082,17 @@ 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)
-      /** The listing's counts read the connector window by window too, never in one statement. */
-      expect(counted).toHaveLength(
+      /**
+       * The tally reads the connector window by window, never in one statement, and keeps the
+       * rows every removal acts on, so no removal reads the connector again.
+       */
+      expect(walked).toEqual([])
+      expect(tallied).toHaveLength(
         Math.floor(
           (present.length + absent.length + tombstoned.length) / RECONCILIATION_WINDOW_SIZE
         ) + 1
       )
-      await expectBounded(counted)
+      await expectBounded(tallied)
     }, 120_000)
 
     it('scans a dense window once and never reads every tombstone of the connector', async () => {
@@ -1109,7 +1112,7 @@ describe('durable source and member cycles in PostgreSQL', () => {
         deletedAt: new Date(startedAt.getTime() - 60_000),
       }))
       await seed([...dense, ...seenTombstones], seenTombstones.length)
-      const { pass, hardDeleted, walked, counted } = await reconcile()
+      const { pass, hardDeleted, walked, tallied } = await reconcile()
       expect(pass).toMatchObject({ complete: true, holdNotice: null })
       expect(hardDeleted.sort()).toEqual(dense.map((item) => item.id).sort())
       expect(
@@ -1119,15 +1122,15 @@ describe('durable source and member cycles in PostgreSQL', () => {
           .where(eq(document.connectorId, connectorId))
       ).toHaveLength(seenTombstones.length)
       /**
-       * The only walk is the hard one: one statement per window, the dense first one included,
-       * then the tail, however many pages of matches a window holds.
+       * One tally statement per window, the dense first one included, then the tail, however
+       * many pages of matches a window holds; the purge pages the rows it kept.
        */
-      expect(walked).toHaveLength(
+      expect(walked).toEqual([])
+      expect(tallied).toHaveLength(
         Math.floor((dense.length + seenTombstones.length) / RECONCILIATION_WINDOW_SIZE) + 1
       )
       /** `deleted_at < $1` implies the tombstone index, which would read every connector tombstone. */
-      await expectBounded(walked)
-      await expectBounded(counted)
+      await expectBounded(tallied)
     }, 120_000)
   })
 

apps/sim/lib/knowledge/connectors/member-sync-engine-content-credential.test.ts

@@ -131,9 +131,9 @@ import {
 } from '@/lib/knowledge/connectors/listing-checkpoint'
 import { executeMemberSync } from '@/lib/knowledge/connectors/member-sync-engine'
 import {
-  routeWindowScans,
-  windowCount,
-  windowCounts,
+  routeWindowTallies,
+  windowTallies,
+  windowTally,
 } from '@/lib/knowledge/connectors/reconciliation-window.test-helpers'
 import {
   MEMBER_SCOPE_RENEWAL_PREFIX_BATCH,
@@ -278,9 +278,7 @@ function arrange(
         : []
     )
     if (!options.contentIncomplete) {
-      windowCounts.push(
-        windowCount({ owned: options.existingDocument ? 2 : 1, listed: 1, soft: 0, hard: 0 })
-      )
+      windowTallies.push(windowTally({ owned: options.existingDocument ? 2 : 1, listed: 1 }))
     }
     if (!options.finishedPass || !options.contentIncomplete)
       queueTableRows(schemaMock.document, [{ count: 1 }])
@@ -367,7 +365,7 @@ describe('member engine with a dedicated content credential', () => {
   beforeEach(() => {
     resetEnvFlagsMock()
     resetDbChainMock()
-    routeWindowScans()
+    routeWindowTallies()
   })
 
   it('refuses an already queued live Search crawl before resolving credentials or taking a lock', async () => {
@@ -855,7 +853,7 @@ describe('member engine with a dedicated content credential', () => {
 
     vi.clearAllMocks()
     resetDbChainMock()
-    routeWindowScans()
+    routeWindowTallies()
     const retry = arrange({
       members: true,
       memberContent: true,

apps/sim/lib/knowledge/connectors/reconciliation-window.test-helpers.ts

@@ -1,25 +1,33 @@
 import { dbChainMockFns } from '@sim/testing/mocks/database.mock'
 import { toRecord } from '@sim/utils/object'
 
-/** Reconciliation window scans, which read through `db.execute`, queued in call order. */
-export const windowScans: unknown[][] = []
+/** Reconciliation window tallies, which read through `db.execute`, queued in call order. */
+export const windowTallies: unknown[][] = []
 
-/** Reconciliation window counts, which read through `db.execute`, queued in call order. */
-export const windowCounts: unknown[][] = []
+interface WindowRow {
+  id: string
+  tombstoned: boolean
+}
 
-/** The one-row result of a completed listing's counts over a connector shorter than a window. */
-export function windowCount(counts: {
+/**
+ * The one-row result of a completed listing's tally over a connector shorter than a window: the
+ * listed count, then the absent documents with ACLs, to soft-delete and to hard-delete.
+ */
+export function windowTally(tally: {
   owned: number
   listed: number
-  acl?: number
-  soft: number
-  hard: number
+  acl?: WindowRow[]
+  soft?: WindowRow[]
+  hard?: WindowRow[]
 }) {
+  const matching: WindowRow[][] = [[], tally.acl ?? [], tally.soft ?? [], tally.hard ?? []]
   return [
     {
-      size: counts.owned,
-      last: counts.owned > 0 ? 'last-owned' : null,
-      matching: [counts.listed, counts.acl ?? 0, counts.soft, counts.hard],
+      size: tally.owned,
+      last: tally.owned > 0 ? 'last-owned' : null,
+      counts: [tally.listed, ...matching.slice(1).map((rows) => rows.length)],
+      ids: matching.map((rows) => rows.map((row) => row.id)),
+      tombstoned: matching.map((rows) => rows.map((row) => row.tombstoned)),
     },
   ]
 }
@@ -43,24 +51,18 @@ function statementText(query: unknown): string {
   return [...strings, ...(Array.isArray(values) ? values.map(statementText) : [])].join(' ')
 }
 
-function windowStatement(query: unknown): 'scan' | 'count' | null {
-  const text = statementText(query)
-  if (!text.includes('MATERIALIZED')) return null
-  return text.includes('json_build_array') ? 'count' : 'scan'
-}
-
 /**
- * Empties {@link windowScans} and {@link windowCounts} and routes `db.execute`: a window scan or
- * count takes the next queued result of its kind, and any other statement reads the database
- * clock, `databaseTime` when given and the worker's clock otherwise.
+ * Empties {@link windowTallies} and routes `db.execute`: a window tally takes the next queued
+ * result, a window scan finds nothing, so a walk no test expects fails on its empty result, and
+ * any other statement reads the database clock, `databaseTime` when given and the worker's clock
+ * otherwise.
  */
-export function routeWindowScans(databaseTime?: Date) {
-  windowScans.length = 0
-  windowCounts.length = 0
+export function routeWindowTallies(databaseTime?: Date) {
+  windowTallies.length = 0
   dbChainMockFns.execute.mockImplementation(async (...args: unknown[]) => {
-    const statement = windowStatement(args[0])
-    if (statement === 'scan') return windowScans.shift() ?? []
-    if (statement === 'count') return windowCounts.shift() ?? []
+    const text = statementText(args[0])
+    if (text.includes('MATERIALIZED'))
+      return text.includes('json_build_array') ? (windowTallies.shift() ?? []) : []
     return [{ startedAt: (databaseTime ?? new Date()).toISOString() }]
   })
 }

apps/sim/lib/knowledge/connectors/reconciliation-window.test.ts

@@ -2,55 +2,104 @@ import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
 import { sql } from 'drizzle-orm'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import {
-  countReconciliationWindows,
+  RECONCILIATION_RETAINED_MATCHES,
   RECONCILIATION_WINDOW_SIZE,
+  tallyReconciliationWindows,
 } from '@/lib/knowledge/connectors/reconciliation-window'
 import {
-  routeWindowScans,
-  windowCounts,
+  routeWindowTallies,
+  windowTallies,
 } from '@/lib/knowledge/connectors/reconciliation-window.test-helpers'
 
-describe('countReconciliationWindows', () => {
-  const conditions = [sql`listed`, sql`absent`]
+describe('tallyReconciliationWindows', () => {
+  const conditions = [
+    { where: sql`listed`, retain: false },
+    { where: sql`absent`, retain: true },
+  ]
+  /** One window's tally row: its size and listed count, then the live absent documents it read. */
+  const window = (size: number, listed: number, absent: string[]) => [
+    {
+      size,
+      last: size > 0 ? `last-${windowTallies.length}` : null,
+      counts: [listed, absent.length],
+      ids: [[], absent],
+      tombstoned: [[], absent.map(() => false)],
+    },
+  ]
+  const absentIds = (prefix: string, count: number) =>
+    Array.from({ length: count }, (_, index) => `${prefix}-${index}`)
+  const tally = (beforePage = async () => undefined) =>
+    tallyReconciliationWindows({
+      connectorId: 'connector',
+      conditions,
+      deadlineAt: Date.now() + 60_000,
+      beforePage,
+    })
 
   beforeEach(() => {
     resetDbChainMock()
-    routeWindowScans()
+    routeWindowTallies()
   })
 
-  it('sums a connector larger than one window across window statements', async () => {
-    windowCounts.push(
-      [{ size: RECONCILIATION_WINDOW_SIZE, last: 'doc-5000', matching: [4000, 7] }],
-      [{ size: 12, last: 'doc-5012', matching: [10, 2] }]
+  it('sums a connector larger than one window and keeps the retained rows of every window', async () => {
+    windowTallies.push(
+      window(RECONCILIATION_WINDOW_SIZE, 4000, ['doc-7', 'doc-9']),
+      window(12, 10, ['doc-5011'])
     )
     const beforePage = vi.fn(async () => undefined)
 
-    const tally = await countReconciliationWindows({
-      connectorId: 'connector',
-      conditions,
-      deadlineAt: Date.now() + 60_000,
-      beforePage,
-    })
-
-    expect(tally).toEqual({
+    expect(await tally(beforePage)).toEqual({
       finished: true,
       owned: RECONCILIATION_WINDOW_SIZE + 12,
-      matching: [4010, 9],
+      matching: [
+        { count: 4010, rows: undefined },
+        {
+          count: 3,
+          rows: ['doc-7', 'doc-9', 'doc-5011'].map((id) => ({ id, tombstoned: false })),
+        },
+      ],
     })
     expect(dbChainMockFns.execute).toHaveBeenCalledTimes(2)
     expect(beforePage).toHaveBeenCalledTimes(2)
-    expect(windowCounts).toEqual([])
+    expect(windowTallies).toEqual([])
+  })
+
+  it('keeps every row up to the retention bound and only the count past it', async () => {
+    const full = RECONCILIATION_WINDOW_SIZE
+    const windows = Math.ceil(RECONCILIATION_RETAINED_MATCHES / full)
+    const fullWindows = () =>
+      Array.from({ length: windows }, (_, index) =>
+        window(full, 0, absentIds(`w${index}`, RECONCILIATION_RETAINED_MATCHES / windows))
+      )
+
+    windowTallies.push(...fullWindows(), window(0, 0, []))
+    const atBound = await tally()
+    expect(atBound.finished && atBound.matching[1]).toEqual({
+      count: RECONCILIATION_RETAINED_MATCHES,
+      rows: Array.from({ length: windows }, (_, index) =>
+        absentIds(`w${index}`, RECONCILIATION_RETAINED_MATCHES / windows)
+      )
+        .flat()
+        .map((id) => ({ id, tombstoned: false })),
+    })
+
+    windowTallies.push(...fullWindows(), window(1, 0, ['one-more']))
+    const pastBound = await tally()
+    expect(pastBound.finished && pastBound.matching[1]).toEqual({
+      count: RECONCILIATION_RETAINED_MATCHES + 1,
+      rows: undefined,
+    })
   })
 
-  it('reports an unfinished count instead of reading past its deadline', async () => {
-    const tally = await countReconciliationWindows({
+  it('reports an unfinished tally instead of reading past its deadline', async () => {
+    const result = await tallyReconciliationWindows({
       connectorId: 'connector',
       conditions,
       deadlineAt: Date.now() - 1,
       beforePage: async () => undefined,
     })
 
-    expect(tally).toEqual({ finished: false })
+    expect(result).toEqual({ finished: false })
     expect(dbChainMockFns.execute).not.toHaveBeenCalled()
   })
 })