Skip to content

Commit 99d034a

Browse files
authored
perf(knowledge): reconcile a completed listing in one pass over its windows (#9449)
The reconciliation count read every window of a connector, then the ACL, soft-delete and hard-delete walks each read the same windows again. The tally now counts every decision input and keeps the rows each removal acts on in the same window statement, so removals page through those rows instead of walking the connector again. A condition matching more than 10,000 documents keeps only its count and falls back to the existing walk, so memory and each statement's result stay bounded.
1 parent 6a269f9 commit 99d034a

9 files changed

Lines changed: 603 additions & 191 deletions

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

Lines changed: 22 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -966,8 +966,8 @@ describe('durable source and member cycles in PostgreSQL', () => {
966966
await db.execute(sql`ANALYZE document`)
967967
}
968968
const isWindowStatement = (query: string) => /^\s*with "document" as materialized/i.test(query)
969-
const isWindowCount = (query: string) => query.includes('json_build_array')
970-
/** Runs the pass, returning the window scans and counts it issued outside transactions. */
969+
const isWindowTally = (query: string) => query.includes('json_build_array')
970+
/** Runs the pass, returning the window walks and tallies it issued outside transactions. */
971971
const reconcile = async () => {
972972
const hardDelete = vi.spyOn(documentService, 'hardDeleteDocuments')
973973
const statements = vi.spyOn(db.$client, 'unsafe')
@@ -1002,10 +1002,10 @@ describe('durable source and member cycles in PostgreSQL', () => {
10021002
hardDeleted: hardDelete.mock.calls.flatMap(([batch]) => batch),
10031003
walked: statements.mock.calls
10041004
.map(([query, params]) => ({ query, params }))
1005-
.filter(({ query }) => isWindowStatement(query) && !isWindowCount(query)),
1006-
counted: statements.mock.calls
1005+
.filter(({ query }) => isWindowStatement(query) && !isWindowTally(query)),
1006+
tallied: statements.mock.calls
10071007
.map(([query, params]) => ({ query, params }))
1008-
.filter(({ query }) => isWindowStatement(query) && isWindowCount(query)),
1008+
.filter(({ query }) => isWindowStatement(query) && isWindowTally(query)),
10091009
}
10101010
} finally {
10111011
statements.mockRestore()
@@ -1042,10 +1042,10 @@ describe('durable source and member cycles in PostgreSQL', () => {
10421042
* must never happen.
10431043
*/
10441044
const expectBounded = async (
1045-
walked: { query: string; params: Parameters<typeof db.$client.unsafe>[1] }[]
1045+
statements: { query: string; params: Parameters<typeof db.$client.unsafe>[1] }[]
10461046
) => {
1047-
expect(walked.length).toBeGreaterThan(0)
1048-
for (const { query, params } of walked) {
1047+
expect(statements.length).toBeGreaterThan(0)
1048+
for (const { query, params } of statements) {
10491049
const { read, scans } = await explainWalk(query, params)
10501050
expect(read, `${query}\n${JSON.stringify(scans, null, 2)}`).toBeLessThanOrEqual(
10511051
RECONCILIATION_WINDOW_SIZE
@@ -1069,7 +1069,7 @@ describe('durable source and member cycles in PostgreSQL', () => {
10691069
deletedAt: new Date(startedAt.getTime() - 60_000),
10701070
}))
10711071
await seed([...present, ...absent, ...tombstoned], present.length)
1072-
const { pass, stats, hardDeleted, walked, counted } = await reconcile()
1072+
const { pass, stats, hardDeleted, walked, tallied } = await reconcile()
10731073
expect(pass).toMatchObject({ complete: true, holdNotice: null })
10741074
expect(stats.docsDeleted).toBe(absent.length)
10751075
expect(hardDeleted.sort()).toEqual(tombstoned.map((item) => item.id).sort())
@@ -1082,14 +1082,17 @@ describe('durable source and member cycles in PostgreSQL', () => {
10821082
expect(
10831083
stored.filter((item) => item.id.includes('-z-')).every((item) => item.acl.length === 0)
10841084
).toBe(true)
1085-
await expectBounded(walked)
1086-
/** The listing's counts read the connector window by window too, never in one statement. */
1087-
expect(counted).toHaveLength(
1085+
/**
1086+
* The tally reads the connector window by window, never in one statement, and keeps the
1087+
* rows every removal acts on, so no removal reads the connector again.
1088+
*/
1089+
expect(walked).toEqual([])
1090+
expect(tallied).toHaveLength(
10881091
Math.floor(
10891092
(present.length + absent.length + tombstoned.length) / RECONCILIATION_WINDOW_SIZE
10901093
) + 1
10911094
)
1092-
await expectBounded(counted)
1095+
await expectBounded(tallied)
10931096
}, 120_000)
10941097

10951098
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', () => {
11091112
deletedAt: new Date(startedAt.getTime() - 60_000),
11101113
}))
11111114
await seed([...dense, ...seenTombstones], seenTombstones.length)
1112-
const { pass, hardDeleted, walked, counted } = await reconcile()
1115+
const { pass, hardDeleted, walked, tallied } = await reconcile()
11131116
expect(pass).toMatchObject({ complete: true, holdNotice: null })
11141117
expect(hardDeleted.sort()).toEqual(dense.map((item) => item.id).sort())
11151118
expect(
@@ -1119,15 +1122,15 @@ describe('durable source and member cycles in PostgreSQL', () => {
11191122
.where(eq(document.connectorId, connectorId))
11201123
).toHaveLength(seenTombstones.length)
11211124
/**
1122-
* The only walk is the hard one: one statement per window, the dense first one included,
1123-
* then the tail, however many pages of matches a window holds.
1125+
* One tally statement per window, the dense first one included, then the tail, however
1126+
* many pages of matches a window holds; the purge pages the rows it kept.
11241127
*/
1125-
expect(walked).toHaveLength(
1128+
expect(walked).toEqual([])
1129+
expect(tallied).toHaveLength(
11261130
Math.floor((dense.length + seenTombstones.length) / RECONCILIATION_WINDOW_SIZE) + 1
11271131
)
11281132
/** `deleted_at < $1` implies the tombstone index, which would read every connector tombstone. */
1129-
await expectBounded(walked)
1130-
await expectBounded(counted)
1133+
await expectBounded(tallied)
11311134
}, 120_000)
11321135
})
11331136

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -131,9 +131,9 @@ import {
131131
} from '@/lib/knowledge/connectors/listing-checkpoint'
132132
import { executeMemberSync } from '@/lib/knowledge/connectors/member-sync-engine'
133133
import {
134-
routeWindowScans,
135-
windowCount,
136-
windowCounts,
134+
routeWindowTallies,
135+
windowTallies,
136+
windowTally,
137137
} from '@/lib/knowledge/connectors/reconciliation-window.test-helpers'
138138
import {
139139
MEMBER_SCOPE_RENEWAL_PREFIX_BATCH,
@@ -278,9 +278,7 @@ function arrange(
278278
: []
279279
)
280280
if (!options.contentIncomplete) {
281-
windowCounts.push(
282-
windowCount({ owned: options.existingDocument ? 2 : 1, listed: 1, soft: 0, hard: 0 })
283-
)
281+
windowTallies.push(windowTally({ owned: options.existingDocument ? 2 : 1, listed: 1 }))
284282
}
285283
if (!options.finishedPass || !options.contentIncomplete)
286284
queueTableRows(schemaMock.document, [{ count: 1 }])
@@ -367,7 +365,7 @@ describe('member engine with a dedicated content credential', () => {
367365
beforeEach(() => {
368366
resetEnvFlagsMock()
369367
resetDbChainMock()
370-
routeWindowScans()
368+
routeWindowTallies()
371369
})
372370

373371
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', () => {
855853

856854
vi.clearAllMocks()
857855
resetDbChainMock()
858-
routeWindowScans()
856+
routeWindowTallies()
859857
const retry = arrange({
860858
members: true,
861859
memberContent: true,
Lines changed: 29 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,33 @@
11
import { dbChainMockFns } from '@sim/testing/mocks/database.mock'
22
import { toRecord } from '@sim/utils/object'
33

4-
/** Reconciliation window scans, which read through `db.execute`, queued in call order. */
5-
export const windowScans: unknown[][] = []
4+
/** Reconciliation window tallies, which read through `db.execute`, queued in call order. */
5+
export const windowTallies: unknown[][] = []
66

7-
/** Reconciliation window counts, which read through `db.execute`, queued in call order. */
8-
export const windowCounts: unknown[][] = []
7+
interface WindowRow {
8+
id: string
9+
tombstoned: boolean
10+
}
911

10-
/** The one-row result of a completed listing's counts over a connector shorter than a window. */
11-
export function windowCount(counts: {
12+
/**
13+
* The one-row result of a completed listing's tally over a connector shorter than a window: the
14+
* listed count, then the absent documents with ACLs, to soft-delete and to hard-delete.
15+
*/
16+
export function windowTally(tally: {
1217
owned: number
1318
listed: number
14-
acl?: number
15-
soft: number
16-
hard: number
19+
acl?: WindowRow[]
20+
soft?: WindowRow[]
21+
hard?: WindowRow[]
1722
}) {
23+
const matching: WindowRow[][] = [[], tally.acl ?? [], tally.soft ?? [], tally.hard ?? []]
1824
return [
1925
{
20-
size: counts.owned,
21-
last: counts.owned > 0 ? 'last-owned' : null,
22-
matching: [counts.listed, counts.acl ?? 0, counts.soft, counts.hard],
26+
size: tally.owned,
27+
last: tally.owned > 0 ? 'last-owned' : null,
28+
counts: [tally.listed, ...matching.slice(1).map((rows) => rows.length)],
29+
ids: matching.map((rows) => rows.map((row) => row.id)),
30+
tombstoned: matching.map((rows) => rows.map((row) => row.tombstoned)),
2331
},
2432
]
2533
}
@@ -43,24 +51,18 @@ function statementText(query: unknown): string {
4351
return [...strings, ...(Array.isArray(values) ? values.map(statementText) : [])].join(' ')
4452
}
4553

46-
function windowStatement(query: unknown): 'scan' | 'count' | null {
47-
const text = statementText(query)
48-
if (!text.includes('MATERIALIZED')) return null
49-
return text.includes('json_build_array') ? 'count' : 'scan'
50-
}
51-
5254
/**
53-
* Empties {@link windowScans} and {@link windowCounts} and routes `db.execute`: a window scan or
54-
* count takes the next queued result of its kind, and any other statement reads the database
55-
* clock, `databaseTime` when given and the worker's clock otherwise.
55+
* Empties {@link windowTallies} and routes `db.execute`: a window tally takes the next queued
56+
* result, a window scan finds nothing, so a walk no test expects fails on its empty result, and
57+
* any other statement reads the database clock, `databaseTime` when given and the worker's clock
58+
* otherwise.
5659
*/
57-
export function routeWindowScans(databaseTime?: Date) {
58-
windowScans.length = 0
59-
windowCounts.length = 0
60+
export function routeWindowTallies(databaseTime?: Date) {
61+
windowTallies.length = 0
6062
dbChainMockFns.execute.mockImplementation(async (...args: unknown[]) => {
61-
const statement = windowStatement(args[0])
62-
if (statement === 'scan') return windowScans.shift() ?? []
63-
if (statement === 'count') return windowCounts.shift() ?? []
63+
const text = statementText(args[0])
64+
if (text.includes('MATERIALIZED'))
65+
return text.includes('json_build_array') ? (windowTallies.shift() ?? []) : []
6466
return [{ startedAt: (databaseTime ?? new Date()).toISOString() }]
6567
})
6668
}

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

Lines changed: 72 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -2,55 +2,104 @@ import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
22
import { sql } from 'drizzle-orm'
33
import { beforeEach, describe, expect, it, vi } from 'vitest'
44
import {
5-
countReconciliationWindows,
5+
RECONCILIATION_RETAINED_MATCHES,
66
RECONCILIATION_WINDOW_SIZE,
7+
tallyReconciliationWindows,
78
} from '@/lib/knowledge/connectors/reconciliation-window'
89
import {
9-
routeWindowScans,
10-
windowCounts,
10+
routeWindowTallies,
11+
windowTallies,
1112
} from '@/lib/knowledge/connectors/reconciliation-window.test-helpers'
1213

13-
describe('countReconciliationWindows', () => {
14-
const conditions = [sql`listed`, sql`absent`]
14+
describe('tallyReconciliationWindows', () => {
15+
const conditions = [
16+
{ where: sql`listed`, retain: false },
17+
{ where: sql`absent`, retain: true },
18+
]
19+
/** One window's tally row: its size and listed count, then the live absent documents it read. */
20+
const window = (size: number, listed: number, absent: string[]) => [
21+
{
22+
size,
23+
last: size > 0 ? `last-${windowTallies.length}` : null,
24+
counts: [listed, absent.length],
25+
ids: [[], absent],
26+
tombstoned: [[], absent.map(() => false)],
27+
},
28+
]
29+
const absentIds = (prefix: string, count: number) =>
30+
Array.from({ length: count }, (_, index) => `${prefix}-${index}`)
31+
const tally = (beforePage = async () => undefined) =>
32+
tallyReconciliationWindows({
33+
connectorId: 'connector',
34+
conditions,
35+
deadlineAt: Date.now() + 60_000,
36+
beforePage,
37+
})
1538

1639
beforeEach(() => {
1740
resetDbChainMock()
18-
routeWindowScans()
41+
routeWindowTallies()
1942
})
2043

21-
it('sums a connector larger than one window across window statements', async () => {
22-
windowCounts.push(
23-
[{ size: RECONCILIATION_WINDOW_SIZE, last: 'doc-5000', matching: [4000, 7] }],
24-
[{ size: 12, last: 'doc-5012', matching: [10, 2] }]
44+
it('sums a connector larger than one window and keeps the retained rows of every window', async () => {
45+
windowTallies.push(
46+
window(RECONCILIATION_WINDOW_SIZE, 4000, ['doc-7', 'doc-9']),
47+
window(12, 10, ['doc-5011'])
2548
)
2649
const beforePage = vi.fn(async () => undefined)
2750

28-
const tally = await countReconciliationWindows({
29-
connectorId: 'connector',
30-
conditions,
31-
deadlineAt: Date.now() + 60_000,
32-
beforePage,
33-
})
34-
35-
expect(tally).toEqual({
51+
expect(await tally(beforePage)).toEqual({
3652
finished: true,
3753
owned: RECONCILIATION_WINDOW_SIZE + 12,
38-
matching: [4010, 9],
54+
matching: [
55+
{ count: 4010, rows: undefined },
56+
{
57+
count: 3,
58+
rows: ['doc-7', 'doc-9', 'doc-5011'].map((id) => ({ id, tombstoned: false })),
59+
},
60+
],
3961
})
4062
expect(dbChainMockFns.execute).toHaveBeenCalledTimes(2)
4163
expect(beforePage).toHaveBeenCalledTimes(2)
42-
expect(windowCounts).toEqual([])
64+
expect(windowTallies).toEqual([])
65+
})
66+
67+
it('keeps every row up to the retention bound and only the count past it', async () => {
68+
const full = RECONCILIATION_WINDOW_SIZE
69+
const windows = Math.ceil(RECONCILIATION_RETAINED_MATCHES / full)
70+
const fullWindows = () =>
71+
Array.from({ length: windows }, (_, index) =>
72+
window(full, 0, absentIds(`w${index}`, RECONCILIATION_RETAINED_MATCHES / windows))
73+
)
74+
75+
windowTallies.push(...fullWindows(), window(0, 0, []))
76+
const atBound = await tally()
77+
expect(atBound.finished && atBound.matching[1]).toEqual({
78+
count: RECONCILIATION_RETAINED_MATCHES,
79+
rows: Array.from({ length: windows }, (_, index) =>
80+
absentIds(`w${index}`, RECONCILIATION_RETAINED_MATCHES / windows)
81+
)
82+
.flat()
83+
.map((id) => ({ id, tombstoned: false })),
84+
})
85+
86+
windowTallies.push(...fullWindows(), window(1, 0, ['one-more']))
87+
const pastBound = await tally()
88+
expect(pastBound.finished && pastBound.matching[1]).toEqual({
89+
count: RECONCILIATION_RETAINED_MATCHES + 1,
90+
rows: undefined,
91+
})
4392
})
4493

45-
it('reports an unfinished count instead of reading past its deadline', async () => {
46-
const tally = await countReconciliationWindows({
94+
it('reports an unfinished tally instead of reading past its deadline', async () => {
95+
const result = await tallyReconciliationWindows({
4796
connectorId: 'connector',
4897
conditions,
4998
deadlineAt: Date.now() - 1,
5099
beforePage: async () => undefined,
51100
})
52101

53-
expect(tally).toEqual({ finished: false })
102+
expect(result).toEqual({ finished: false })
54103
expect(dbChainMockFns.execute).not.toHaveBeenCalled()
55104
})
56105
})

0 commit comments

Comments
 (0)