improvement(knowledge): bound the projection warm and contain its extension probe

Author: waleedlatif1

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

@@ -95,6 +95,79 @@ describe('prewarmSearchProjection', () => {
     expect(warmed.map((item) => item.relation)).toEqual(['embedding_search_512_cosine_hnsw_idx'])
   })
 
+  it('returns nothing when the extension cannot be checked, never failing its caller', async () => {
+    const fake = session({ installed: true })
+    fake.unsafe = async () => {
+      throw new Error('canceling statement due to user request')
+    }
+    await expect(prewarmSearchProjection(fake)).resolves.toEqual([])
+  })
+
+  it('bounds every read by the budget left and leaves the rest cold once it is spent', async () => {
+    vi.useFakeTimers()
+    try {
+      const fake = session({
+        installed: true,
+        relations: [
+          'embedding_search',
+          'embedding_keyword_tin',
+          'embedding_search_512_cosine_hnsw_idx',
+        ],
+      })
+      const read = fake.unsafe
+      fake.unsafe = async (query: string, parameters?: string[]) => {
+        const rows = await read(query, parameters)
+        /** Each read takes 400 ms of a 1 s budget. */
+        if (query.includes('pg_prewarm(')) vi.advanceTimersByTime(400)
+        return rows
+      }
+      const warmed = await prewarmSearchProjection(fake, { budgetMs: 1000 })
+      expect(warmed.map((item) => item.relation)).toEqual([
+        'embedding_search',
+        'embedding_keyword_tin',
+        'embedding_search_512_cosine_hnsw_idx',
+      ])
+      const timeouts = fake.statements
+        .filter((statement) => statement.query.startsWith('SET statement_timeout'))
+        .map((statement) => Number(statement.query.split('= ')[1]))
+      expect(timeouts).toEqual([1000, 600, 200])
+      expect(fake.statements.at(-1)?.query).toBe('RESET statement_timeout')
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
+  it('skips the relations beyond a spent budget', async () => {
+    vi.useFakeTimers()
+    try {
+      const fake = session({
+        installed: true,
+        relations: ['embedding_search', 'embedding_search_512_cosine_hnsw_idx'],
+      })
+      const read = fake.unsafe
+      fake.unsafe = async (query: string, parameters?: string[]) => {
+        const rows = await read(query, parameters)
+        if (query.includes('pg_prewarm(')) vi.advanceTimersByTime(1500)
+        return rows
+      }
+      const warmed = await prewarmSearchProjection(fake, { budgetMs: 1000 })
+      expect(warmed.map((item) => item.relation)).toEqual(['embedding_search'])
+      expect(
+        fake.statements.filter((statement) => statement.query.includes('pg_prewarm('))
+      ).toHaveLength(1)
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
+  it('never sets a timeout on an unbounded pass', async () => {
+    const fake = session({ installed: true, relations: ['embedding_search'] })
+    await prewarmSearchProjection(fake)
+    expect(fake.statements.some((statement) => statement.query.includes('statement_timeout'))).toBe(
+      false
+    )
+  })
+
   it('returns nothing when the catalog cannot be read, never failing its caller', async () => {
     const fake = session({ installed: true })
     fake.unsafe = async (query: string) => {

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

@@ -22,6 +22,15 @@ export interface PrewarmedRelation {
   elapsedMs: number
 }
 
+export interface PrewarmOptions {
+  /**
+   * Wall-clock ceiling for the whole pass. Each read is bounded by the time left, and relations
+   * beyond the ceiling stay cold; a caller with its own run limit sets it so warming can never
+   * outlive the run that asked for it.
+   */
+  budgetMs?: number
+}
+
 /**
  * `pg_prewarm` is not a trusted extension, so the application role cannot create it and no
  * migration can; a superuser installs it once. Without it the projection warms only as searches
@@ -57,18 +66,23 @@ export async function prewarmRelation(
  * Heaps go first and the ranking indexes last, so where the cache cannot hold everything the
  * indexes are what survives: a walk reads far more index pages than heap pages. Relations are
  * resolved through the search path, so a schema that carries its own copy warms its own copy.
- * A relation that fails to warm is logged and skipped; warming is never worth failing the
+ * Nothing here throws: a missing extension, an unreadable catalog, a relation that fails to
+ * read or a spent budget is logged and skipped, since warming is never worth failing the
  * operation that asked for it.
  */
 export async function prewarmSearchProjection(
-  session: PrewarmSession
+  session: PrewarmSession,
+  options: PrewarmOptions = {}
 ): Promise<PrewarmedRelation[]> {
-  if (!(await pgPrewarmInstalled(session))) {
-    logger.warn('pg_prewarm is not installed; the search projection warms only as it is searched')
-    return []
-  }
+  const startedAt = Date.now()
+  const remainingMs = () =>
+    options.budgetMs === undefined ? undefined : options.budgetMs - (Date.now() - startedAt)
   let relations: string[]
   try {
+    if (!(await pgPrewarmInstalled(session))) {
+      logger.warn('pg_prewarm is not installed; the search projection warms only as it is searched')
+      return []
+    }
     relations = await rankingRelations(session)
   } catch (error) {
     logger.warn('Search projection relations could not be listed', {
@@ -77,20 +91,37 @@ export async function prewarmSearchProjection(
     return []
   }
   const warmed: PrewarmedRelation[] = []
-  for (const relation of relations) {
-    try {
-      warmed.push(await prewarmRelation(session, relation))
-    } catch (error) {
-      logger.warn('Search projection relation failed to warm', {
-        relation,
-        error: getErrorMessage(error),
-      })
+  const cold: string[] = []
+  try {
+    for (const relation of relations) {
+      const left = remainingMs()
+      if (left !== undefined && left <= 0) {
+        cold.push(relation)
+        continue
+      }
+      try {
+        if (left !== undefined) {
+          await session.unsafe(`SET statement_timeout = ${Math.ceil(left)}`)
+        }
+        warmed.push(await prewarmRelation(session, relation))
+      } catch (error) {
+        cold.push(relation)
+        logger.warn('Search projection relation failed to warm', {
+          relation,
+          error: getErrorMessage(error),
+        })
+      }
+    }
+  } finally {
+    if (options.budgetMs !== undefined) {
+      await Promise.resolve(session.unsafe('RESET statement_timeout')).catch(() => undefined)
     }
   }
   logger.info('Search projection warmed', {
     relations: warmed.length,
+    cold,
     pages: warmed.reduce((sum, item) => sum + item.pages, 0),
-    elapsedMs: warmed.reduce((sum, item) => sum + item.elapsedMs, 0),
+    elapsedMs: Date.now() - startedAt,
   })
   return warmed
 }

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

@@ -28,6 +28,7 @@ vi.mock('@/lib/core/utils/background', () => ({
 
 import {
   enqueueProjectionSourceAclBackfill,
+  PROJECTION_PREWARM_BUDGET_MS,
   runProjectionSourceAclBackfill,
 } from '@/lib/knowledge/search/projection-source-acl-backfill'
 
@@ -62,7 +63,7 @@ describe('runProjectionSourceAclBackfill', () => {
   it('warms the projections on the same connection once both are filled, before closing it', async () => {
     await runProjectionSourceAclBackfill({})
     expect(mockPrewarm).toHaveBeenCalledTimes(1)
-    expect(mockPrewarm).toHaveBeenCalledWith(connection)
+    expect(mockPrewarm).toHaveBeenCalledWith(connection, { budgetMs: PROJECTION_PREWARM_BUDGET_MS })
     expect(mockPrewarm.mock.invocationCallOrder[0]).toBeLessThan(
       mockEnd.mock.invocationCallOrder[0]
     )

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

@@ -16,6 +16,13 @@ const logger = createLogger('ProjectionSourceAclBackfill')
 
 export const PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID = 'projection-source-acl-backfill'
 
+/**
+ * Ceiling on warming the projections after the fill. It sits inside the headroom the worker
+ * keeps beyond a run's fill budget, so a slow read of a large projection can never carry the
+ * completed run past the worker's limit and repeat the fill on retry.
+ */
+export const PROJECTION_PREWARM_BUDGET_MS = 15 * 60 * 1000
+
 /** Where a run stopped, so the next one carries on from there instead of rescanning. */
 export interface ProjectionSourceAclBackfillCursor {
   projection: ProjectionSourceAclTable
@@ -71,7 +78,7 @@ export async function runProjectionSourceAclBackfill(
       elapsedMs: Date.now() - startedAt,
     })
     /** The fill just streamed through both projections; put the ranking pages back before anyone searches. */
-    await prewarmSearchProjection(sql)
+    await prewarmSearchProjection(sql, { budgetMs: PROJECTION_PREWARM_BUDGET_MS })
     return null
   } finally {
     await sql.end()