Skip to content

Commit 5457698

Browse files
waleedlatif1waleedlatif1
authored
improvement(knowledge): warm the search projection after its backfill (#8092)
* improvement(knowledge): warm the search projection after its backfill * improvement(knowledge): bound the projection warm and contain its extension probe
1 parent ea31ca7 commit 5457698

5 files changed

Lines changed: 410 additions & 1 deletion

File tree

‎apps/sim/lib/knowledge/search/prewarm.test.ts‎

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
/**
2+
* @vitest-environment node
3+
*/
4+
import { beforeEach, describe, expect, it, vi } from 'vitest'
5+
6+
vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({
7+
PROJECTION_SOURCE_ACL_TABLES: ['embedding_search', 'embedding_keyword_tin'],
8+
}))
9+
10+
import {
11+
pgPrewarmInstalled,
12+
prewarmRelation,
13+
prewarmSearchProjection,
14+
} from '@/lib/knowledge/search/prewarm'
15+
16+
interface Statement {
17+
query: string
18+
parameters?: string[]
19+
}
20+
21+
/** A session that records every statement and answers from the case's catalog. */
22+
function session(state: { installed: boolean; relations?: string[]; failing?: string[] }): {
23+
statements: Statement[]
24+
unsafe: (query: string, parameters?: string[]) => Promise<unknown[]>
25+
} {
26+
const statements: Statement[] = []
27+
return {
28+
statements,
29+
unsafe: async (query: string, parameters?: string[]) => {
30+
statements.push({ query, parameters })
31+
if (query.includes('pg_extension')) return state.installed ? [{ '?column?': 1 }] : []
32+
if (query.includes('pg_class'))
33+
return (state.relations ?? []).map((relation) => ({ relation }))
34+
if (query.includes('pg_prewarm(')) {
35+
const [relation] = parameters ?? []
36+
if (state.failing?.includes(relation))
37+
throw new Error(`relation "${relation}" does not exist`)
38+
return [{ pages: 7 }]
39+
}
40+
return []
41+
},
42+
}
43+
}
44+
45+
describe('prewarmSearchProjection', () => {
46+
beforeEach(() => {
47+
vi.clearAllMocks()
48+
})
49+
50+
it('does nothing where the extension is absent, so the application role never needs it', async () => {
51+
const fake = session({ installed: false })
52+
await expect(prewarmSearchProjection(fake)).resolves.toEqual([])
53+
expect(fake.statements).toHaveLength(1)
54+
expect(fake.statements[0].query).toContain("extname = 'pg_prewarm'")
55+
})
56+
57+
it('reads the projections and their ranking indexes in the order the catalog lists them', async () => {
58+
const fake = session({
59+
installed: true,
60+
relations: [
61+
'embedding_search',
62+
'embedding_keyword_tin',
63+
'embedding_search_512_cosine_hnsw_idx',
64+
],
65+
})
66+
const warmed = await prewarmSearchProjection(fake)
67+
expect(warmed.map((item) => item.relation)).toEqual([
68+
'embedding_search',
69+
'embedding_keyword_tin',
70+
'embedding_search_512_cosine_hnsw_idx',
71+
])
72+
expect(warmed.every((item) => item.pages === 7)).toBe(true)
73+
const listed = fake.statements.find((statement) => statement.query.includes('pg_class'))
74+
expect(listed?.parameters).toEqual([
75+
'{embedding_search,embedding_keyword_tin}',
76+
'{hnsw,tin,gin}',
77+
])
78+
expect(listed?.query).toContain("ORDER BY c.relkind = 'r' DESC")
79+
const reads = fake.statements.filter((statement) => statement.query.includes('pg_prewarm('))
80+
expect(reads.map((statement) => statement.parameters)).toEqual([
81+
['embedding_search'],
82+
['embedding_keyword_tin'],
83+
['embedding_search_512_cosine_hnsw_idx'],
84+
])
85+
expect(reads.every((statement) => statement.query.includes("'read'"))).toBe(true)
86+
})
87+
88+
it('skips a relation that fails to warm and carries on with the rest', async () => {
89+
const fake = session({
90+
installed: true,
91+
relations: ['embedding_search', 'embedding_search_512_cosine_hnsw_idx'],
92+
failing: ['embedding_search'],
93+
})
94+
const warmed = await prewarmSearchProjection(fake)
95+
expect(warmed.map((item) => item.relation)).toEqual(['embedding_search_512_cosine_hnsw_idx'])
96+
})
97+
98+
it('returns nothing when the extension cannot be checked, never failing its caller', async () => {
99+
const fake = session({ installed: true })
100+
fake.unsafe = async () => {
101+
throw new Error('canceling statement due to user request')
102+
}
103+
await expect(prewarmSearchProjection(fake)).resolves.toEqual([])
104+
})
105+
106+
it('bounds every read by the budget left and leaves the rest cold once it is spent', async () => {
107+
vi.useFakeTimers()
108+
try {
109+
const fake = session({
110+
installed: true,
111+
relations: [
112+
'embedding_search',
113+
'embedding_keyword_tin',
114+
'embedding_search_512_cosine_hnsw_idx',
115+
],
116+
})
117+
const read = fake.unsafe
118+
fake.unsafe = async (query: string, parameters?: string[]) => {
119+
const rows = await read(query, parameters)
120+
/** Each read takes 400 ms of a 1 s budget. */
121+
if (query.includes('pg_prewarm(')) vi.advanceTimersByTime(400)
122+
return rows
123+
}
124+
const warmed = await prewarmSearchProjection(fake, { budgetMs: 1000 })
125+
expect(warmed.map((item) => item.relation)).toEqual([
126+
'embedding_search',
127+
'embedding_keyword_tin',
128+
'embedding_search_512_cosine_hnsw_idx',
129+
])
130+
const timeouts = fake.statements
131+
.filter((statement) => statement.query.startsWith('SET statement_timeout'))
132+
.map((statement) => Number(statement.query.split('= ')[1]))
133+
expect(timeouts).toEqual([1000, 600, 200])
134+
expect(fake.statements.at(-1)?.query).toBe('RESET statement_timeout')
135+
} finally {
136+
vi.useRealTimers()
137+
}
138+
})
139+
140+
it('skips the relations beyond a spent budget', async () => {
141+
vi.useFakeTimers()
142+
try {
143+
const fake = session({
144+
installed: true,
145+
relations: ['embedding_search', 'embedding_search_512_cosine_hnsw_idx'],
146+
})
147+
const read = fake.unsafe
148+
fake.unsafe = async (query: string, parameters?: string[]) => {
149+
const rows = await read(query, parameters)
150+
if (query.includes('pg_prewarm(')) vi.advanceTimersByTime(1500)
151+
return rows
152+
}
153+
const warmed = await prewarmSearchProjection(fake, { budgetMs: 1000 })
154+
expect(warmed.map((item) => item.relation)).toEqual(['embedding_search'])
155+
expect(
156+
fake.statements.filter((statement) => statement.query.includes('pg_prewarm('))
157+
).toHaveLength(1)
158+
} finally {
159+
vi.useRealTimers()
160+
}
161+
})
162+
163+
it('never sets a timeout on an unbounded pass', async () => {
164+
const fake = session({ installed: true, relations: ['embedding_search'] })
165+
await prewarmSearchProjection(fake)
166+
expect(fake.statements.some((statement) => statement.query.includes('statement_timeout'))).toBe(
167+
false
168+
)
169+
})
170+
171+
it('returns nothing when the catalog cannot be read, never failing its caller', async () => {
172+
const fake = session({ installed: true })
173+
fake.unsafe = async (query: string) => {
174+
if (query.includes('pg_extension')) return [{ '?column?': 1 }]
175+
throw new Error('permission denied for table pg_class')
176+
}
177+
await expect(prewarmSearchProjection(fake)).resolves.toEqual([])
178+
})
179+
})
180+
181+
describe('prewarmRelation', () => {
182+
it('reports the pages read for one relation', async () => {
183+
const fake = session({ installed: true })
184+
await expect(prewarmRelation(fake, 'embedding_search')).resolves.toMatchObject({
185+
relation: 'embedding_search',
186+
pages: 7,
187+
})
188+
expect(await pgPrewarmInstalled(fake)).toBe(true)
189+
})
190+
})

‎apps/sim/lib/knowledge/search/prewarm.ts‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
import { PROJECTION_SOURCE_ACL_TABLES } from '@sim/db/script-migrations/0021_embedding_search_connector'
2+
import { createLogger } from '@sim/logger'
3+
import { getErrorMessage } from '@sim/utils/errors'
4+
5+
const logger = createLogger('SearchProjectionPrewarm')
6+
7+
/**
8+
* The access methods a ranking touches at random: the vector graphs, the Tin keyword index, and
9+
* the GIN index the on-row permission test reads. The remaining b-trees serve hydration, which
10+
* reads a handful of rows by key and is fast cold.
11+
*/
12+
const RANKING_ACCESS_METHODS = ['hnsw', 'tin', 'gin'] as const
13+
14+
/** The one call the helper needs from a `postgres` connection or a reserved session. */
15+
export interface PrewarmSession {
16+
unsafe(query: string, parameters?: string[]): PromiseLike<ArrayLike<Record<string, unknown>>>
17+
}
18+
19+
export interface PrewarmedRelation {
20+
relation: string
21+
pages: number
22+
elapsedMs: number
23+
}
24+
25+
export interface PrewarmOptions {
26+
/**
27+
* Wall-clock ceiling for the whole pass. Each read is bounded by the time left, and relations
28+
* beyond the ceiling stay cold; a caller with its own run limit sets it so warming can never
29+
* outlive the run that asked for it.
30+
*/
31+
budgetMs?: number
32+
}
33+
34+
/**
35+
* `pg_prewarm` is not a trusted extension, so the application role cannot create it and no
36+
* migration can; a superuser installs it once. Without it the projection warms only as searches
37+
* touch it, which is what a bulk operation leaves behind.
38+
*/
39+
export async function pgPrewarmInstalled(session: PrewarmSession): Promise<boolean> {
40+
const rows = await session.unsafe("SELECT 1 FROM pg_extension WHERE extname = 'pg_prewarm'")
41+
return rows.length > 0
42+
}
43+
44+
/**
45+
* Reads one relation into the operating system's cache. `read` mode leaves shared buffers to the
46+
* workload, where `buffer` mode would evict them wholesale to make room.
47+
*/
48+
export async function prewarmRelation(
49+
session: PrewarmSession,
50+
relation: string
51+
): Promise<PrewarmedRelation> {
52+
const startedAt = Date.now()
53+
const [row] = Array.from(
54+
await session.unsafe("SELECT pg_prewarm($1::regclass, 'read')::int AS pages", [relation])
55+
)
56+
return { relation, pages: Number(row?.pages ?? 0), elapsedMs: Date.now() - startedAt }
57+
}
58+
59+
/**
60+
* Warms the ranking projections after something streamed through them. A backfill or index build
61+
* reads every heap page in order and pushes the vector graphs out of cache; the next searches
62+
* then fetch the graph one random page at a time from disk, take seconds, and end at their
63+
* deadline with partial results. Reading the projections back in makes the first search after a
64+
* bulk operation as fast as the thousandth.
65+
*
66+
* Heaps go first and the ranking indexes last, so where the cache cannot hold everything the
67+
* indexes are what survives: a walk reads far more index pages than heap pages. Relations are
68+
* resolved through the search path, so a schema that carries its own copy warms its own copy.
69+
* Nothing here throws: a missing extension, an unreadable catalog, a relation that fails to
70+
* read or a spent budget is logged and skipped, since warming is never worth failing the
71+
* operation that asked for it.
72+
*/
73+
export async function prewarmSearchProjection(
74+
session: PrewarmSession,
75+
options: PrewarmOptions = {}
76+
): Promise<PrewarmedRelation[]> {
77+
const startedAt = Date.now()
78+
const remainingMs = () =>
79+
options.budgetMs === undefined ? undefined : options.budgetMs - (Date.now() - startedAt)
80+
let relations: string[]
81+
try {
82+
if (!(await pgPrewarmInstalled(session))) {
83+
logger.warn('pg_prewarm is not installed; the search projection warms only as it is searched')
84+
return []
85+
}
86+
relations = await rankingRelations(session)
87+
} catch (error) {
88+
logger.warn('Search projection relations could not be listed', {
89+
error: getErrorMessage(error),
90+
})
91+
return []
92+
}
93+
const warmed: PrewarmedRelation[] = []
94+
const cold: string[] = []
95+
try {
96+
for (const relation of relations) {
97+
const left = remainingMs()
98+
if (left !== undefined && left <= 0) {
99+
cold.push(relation)
100+
continue
101+
}
102+
try {
103+
if (left !== undefined) {
104+
await session.unsafe(`SET statement_timeout = ${Math.ceil(left)}`)
105+
}
106+
warmed.push(await prewarmRelation(session, relation))
107+
} catch (error) {
108+
cold.push(relation)
109+
logger.warn('Search projection relation failed to warm', {
110+
relation,
111+
error: getErrorMessage(error),
112+
})
113+
}
114+
}
115+
} finally {
116+
if (options.budgetMs !== undefined) {
117+
await Promise.resolve(session.unsafe('RESET statement_timeout')).catch(() => undefined)
118+
}
119+
}
120+
logger.info('Search projection warmed', {
121+
relations: warmed.length,
122+
cold,
123+
pages: warmed.reduce((sum, item) => sum + item.pages, 0),
124+
elapsedMs: Date.now() - startedAt,
125+
})
126+
return warmed
127+
}
128+
129+
/** The projections' heaps, then their ranking indexes smallest first, as the search path finds them. */
130+
async function rankingRelations(session: PrewarmSession): Promise<string[]> {
131+
const rows = await session.unsafe(
132+
`WITH heaps AS (
133+
SELECT to_regclass(name) AS oid FROM unnest($1::text[]) AS name
134+
)
135+
SELECT c.oid::regclass::text AS relation
136+
FROM pg_class c
137+
JOIN pg_am am ON am.oid = c.relam
138+
LEFT JOIN pg_index i ON i.indexrelid = c.oid
139+
WHERE c.oid IN (SELECT oid FROM heaps)
140+
OR (
141+
i.indrelid IN (SELECT oid FROM heaps)
142+
AND i.indisvalid
143+
AND am.amname = ANY($2::text[])
144+
)
145+
ORDER BY c.relkind = 'r' DESC, pg_relation_size(c.oid)`,
146+
[toArrayLiteral(PROJECTION_SOURCE_ACL_TABLES), toArrayLiteral(RANKING_ACCESS_METHODS)]
147+
)
148+
return Array.from(rows, (row) => String(row.relation))
149+
}
150+
151+
/** Postgres array literal for identifiers that carry no quotes, commas or braces. */
152+
function toArrayLiteral(values: readonly string[]): string {
153+
return `{${values.join(',')}}`
154+
}

‎apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,11 @@
33
*/
44
import { beforeEach, describe, expect, it, vi } from 'vitest'
55

6-
const { mockBackfill, mockEnd, mockPostgres, mockTasksTrigger } = vi.hoisted(() => ({
6+
const { mockBackfill, mockEnd, mockPostgres, mockPrewarm, mockTasksTrigger } = vi.hoisted(() => ({
77
mockBackfill: vi.fn(),
88
mockEnd: vi.fn(async () => undefined),
99
mockPostgres: vi.fn(),
10+
mockPrewarm: vi.fn(async () => []),
1011
mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })),
1112
}))
1213

@@ -16,6 +17,7 @@ vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({
1617
backfillProjectionSourceAcl: mockBackfill,
1718
}))
1819
vi.mock('postgres', () => ({ default: mockPostgres }))
20+
vi.mock('@/lib/knowledge/search/prewarm', () => ({ prewarmSearchProjection: mockPrewarm }))
1921
vi.mock('@trigger.dev/sdk', () => ({ tasks: { trigger: mockTasksTrigger } }))
2022
vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' }))
2123
vi.mock('@/lib/core/utils/background', () => ({
@@ -26,6 +28,7 @@ vi.mock('@/lib/core/utils/background', () => ({
2628

2729
import {
2830
enqueueProjectionSourceAclBackfill,
31+
PROJECTION_PREWARM_BUDGET_MS,
2932
runProjectionSourceAclBackfill,
3033
} from '@/lib/knowledge/search/projection-source-acl-backfill'
3134

@@ -57,6 +60,15 @@ describe('runProjectionSourceAclBackfill', () => {
5760
expect(mockEnd).toHaveBeenCalledTimes(1)
5861
})
5962

63+
it('warms the projections on the same connection once both are filled, before closing it', async () => {
64+
await runProjectionSourceAclBackfill({})
65+
expect(mockPrewarm).toHaveBeenCalledTimes(1)
66+
expect(mockPrewarm).toHaveBeenCalledWith(connection, { budgetMs: PROJECTION_PREWARM_BUDGET_MS })
67+
expect(mockPrewarm.mock.invocationCallOrder[0]).toBeLessThan(
68+
mockEnd.mock.invocationCallOrder[0]
69+
)
70+
})
71+
6072
it('resumes after the cursor in its projection and from the start of the next', async () => {
6173
await runProjectionSourceAclBackfill({
6274
cursor: { projection: 'embedding_keyword_tin', afterId: 'chunk-9' },
@@ -80,6 +92,7 @@ describe('runProjectionSourceAclBackfill', () => {
8092
})
8193
expect(mockBackfill).toHaveBeenCalledTimes(1)
8294
expect(mockBackfill.mock.calls[0][2].budgetMs).toBeLessThanOrEqual(1000)
95+
expect(mockPrewarm).not.toHaveBeenCalled()
8396
expect(mockEnd).toHaveBeenCalledTimes(1)
8497
})
8598

0 commit comments

Comments
 (0)