11import { db } from '@sim/db'
22import { document } from '@sim/db/schema'
3- import { and , asc , eq , gt , isNull , lte , type SQL , sql } from 'drizzle-orm'
3+ import { and , asc , eq , gt , isNull , type SQL , sql } from 'drizzle-orm'
44
55/**
66 * Ids of `doc_connector_reconciliation_v2_idx` one reconciliation statement may scan. The rows a
77 * walk looks for can be rare and late in id order, so a page bounded only by its matches could
8- * read, and heap-fetch, a whole connector in one statement. Each page instead scans at most this
9- * many ids: a few hundred milliseconds cold, and few enough round trips that a large connector
10- * is walked in hundreds of pages rather than thousands.
8+ * read, and heap-fetch, a whole connector in one statement. Each window instead scans at most
9+ * this many ids: a few hundred milliseconds cold, and few enough round trips that a large
10+ * connector is walked in hundreds of statements rather than thousands.
1111 */
1212export const RECONCILIATION_WINDOW_SIZE = 5_000
1313
14- /** A type alias, not an interface, so it satisfies `db.execute`'s row-record constraint. */
15- type ReconciliationRow = {
14+ interface ReconciliationRow {
1615 id : string
1716 tombstoned : boolean
1817}
@@ -31,40 +30,27 @@ export interface ReconciliationWalk {
3130 onPage : ( rows : ReconciliationRow [ ] ) => Promise < void >
3231}
3332
34- function ownedAfter ( connectorId : string , afterId : string | undefined ) {
35- return and (
36- eq ( document . connectorId , connectorId ) ,
37- eq ( document . userExcluded , false ) ,
38- isNull ( document . archivedAt ) ,
39- afterId ? gt ( document . id , afterId ) : undefined
40- )
41- }
42-
43- /** The last id of the next full window after `afterId`, or undefined when less than a window remains. */
44- async function windowEnd ( connectorId : string , afterId : string | undefined ) {
45- const [ row ] = await db
46- . select ( { id : document . id } )
47- . from ( document )
48- . where ( ownedAfter ( connectorId , afterId ) )
49- . orderBy ( asc ( document . id ) )
50- . limit ( 1 )
51- . offset ( RECONCILIATION_WINDOW_SIZE - 1 )
52- return row ?. id
33+ /** A type alias, not an interface, so it satisfies `db.execute`'s row-record constraint. */
34+ type WindowScan = {
35+ size : number
36+ last : string | null
37+ ids : string [ ]
38+ tombstoned : boolean [ ]
5339}
5440
5541/**
56- * The first `pageSize` rows of the window `(afterId, endId]` matching `condition`. The window is
57- * read on its own, as a materialized CTE that shadows `document`, so the planner reaches it only
58- * through the v2 id range: a condition that implies a narrower partial index, such as the
59- * tombstone index for `deleted_at < $1`, would otherwise draw a bitmap over every tombstone the
60- * connector has.
42+ * Reads the next window after `afterId` and the ids in it matching `condition`, in one statement.
43+ * The matches come back as JSON, which the driver decodes without the array type lookup.
44+ * The window is the first {@link RECONCILIATION_WINDOW_SIZE} owned documents in id order, a
45+ * materialized CTE that shadows `document`: the `ORDER BY id LIMIT` makes it an ordered walk of
46+ * the v2 index that stops at the window's size. A range bounded only by `id` is not enough, as
47+ * the planner cannot tell that it covers one window of the connector and may instead read every
48+ * document of the connector through another connector index; and a condition that implies a
49+ * narrower partial index, such as the tombstone index for `deleted_at < $1`, could draw a bitmap
50+ * over every tombstone the connector has.
6151 */
62- async function windowPage (
63- walk : ReconciliationWalk ,
64- afterId : string | undefined ,
65- endId : string | undefined
66- ) : Promise < ReconciliationRow [ ] > {
67- const rows = db
52+ async function scanWindow ( walk : ReconciliationWalk , afterId : string | undefined ) {
53+ const window = db
6854 . select ( {
6955 id : document . id ,
7056 connectorId : document . connectorId ,
@@ -76,42 +62,56 @@ async function windowPage(
7662 contentHash : document . contentHash ,
7763 } )
7864 . from ( document )
79- . where ( and ( ownedAfter ( walk . connectorId , afterId ) , endId ? lte ( document . id , endId ) : undefined ) )
80- return db . execute < ReconciliationRow > ( sql `
81- WITH ${ document } ${ rows }
82- SELECT ${ document . id } ${ document . deletedAt }
83- FROM ${ document }
84- WHERE ${ walk . condition ?? sql `true` }
85- ORDER BY ${ document . id }
86- LIMIT ${ walk . pageSize } )
65+ . where (
66+ and (
67+ eq ( document . connectorId , walk . connectorId ) ,
68+ eq ( document . userExcluded , false ) ,
69+ isNull ( document . archivedAt ) ,
70+ afterId ? gt ( document . id , afterId ) : undefined
71+ )
72+ )
73+ . orderBy ( asc ( document . id ) )
74+ . limit ( RECONCILIATION_WINDOW_SIZE )
75+ const [ scan ] = await db . execute < WindowScan > ( sql `
76+ WITH ${ document } ${ window }
77+ matched AS (
78+ SELECT ${ document . id } ${ document . deletedAt }
79+ FROM ${ document }
80+ WHERE ${ walk . condition ?? sql `true` }
81+ )
82+ SELECT
83+ (SELECT count(*)::int FROM ${ document }
84+ (SELECT max(${ document . id } ${ document }
85+ coalesce(json_agg(id ORDER BY id), '[]') AS "ids",
86+ coalesce(json_agg(tombstoned ORDER BY id), '[]') AS "tombstoned"
87+ FROM matched` )
88+ return scan
8789}
8890
8991/**
90- * Walks a connector's owned documents in id order, one bounded window per statement, handing
91- * each page of the rows matching `condition` to `onPage`. A window is bounded once and paged
92- * until its matches run out, so a dense window costs one bound however many pages it holds. A
93- * window without matches still advances the walk, which ends only when the connector's tail (a
94- * window with no end) is exhausted. Returns false when the deadline stopped it.
92+ * Walks a connector's owned documents in id order, one window per statement, handing the rows
93+ * matching `condition` to `onPage` in pages of at most `pageSize`. A window without matches still
94+ * advances the walk, which ends after the first window shorter than a full one. Returns false when
95+ * the deadline stopped it.
9596 */
9697export async function walkReconciliationWindows ( walk : ReconciliationWalk ) : Promise < boolean > {
9798 let after : string | undefined
98- let window : { end : string | undefined } | undefined
9999 for ( ; ; ) {
100100 if ( Date . now ( ) >= walk . deadlineAt ) return false
101101 await walk . beforePage ( )
102102 if ( Date . now ( ) >= walk . deadlineAt ) return false
103- window ??= { end : await windowEnd ( walk . connectorId , after ) }
104- const { end } = window
105- const rows = await windowPage ( walk , after , end )
106- if ( rows . length > 0 ) {
103+ const { size, last, ids, tombstoned } = await scanWindow ( walk , after )
104+ for ( let offset = 0 ; offset < ids . length ; offset += walk . pageSize ) {
105+ if ( offset > 0 ) await walk . beforePage ( )
107106 /** Materialized ids are acted on only inside the budget, so a late page is left for the next run. */
108107 if ( Date . now ( ) >= walk . deadlineAt ) return false
109- await walk . onPage ( rows )
108+ await walk . onPage (
109+ ids
110+ . slice ( offset , offset + walk . pageSize )
111+ . map ( ( id , index ) => ( { id, tombstoned : tombstoned [ offset + index ] } ) )
112+ )
110113 }
111- if ( rows . length === walk . pageSize ) after = rows [ rows . length - 1 ] . id
112- else if ( end ) {
113- after = end
114- window = undefined
115- } else return true
114+ if ( size < RECONCILIATION_WINDOW_SIZE || ! last ) return true
115+ after = last
116116 }
117117}
0 commit comments