improvement(knowledge): warm the search projection after its backfill

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

@@ -0,0 +1,190 @@
+/**
+ * @vitest-environment node
+ */
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({
+  PROJECTION_SOURCE_ACL_TABLES: ['embedding_search', 'embedding_keyword_tin'],
+}))
+
+import {
+  pgPrewarmInstalled,
+  prewarmRelation,
+  prewarmSearchProjection,
+} from '@/lib/knowledge/search/prewarm'
+
+interface Statement {
+  query: string
+  parameters?: string[]
+}
+
+/** A session that records every statement and answers from the case's catalog. */
+function session(state: { installed: boolean; relations?: string[]; failing?: string[] }): {
+  statements: Statement[]
+  unsafe: (query: string, parameters?: string[]) => Promise<unknown[]>
+} {
+  const statements: Statement[] = []
+  return {
+    statements,
+    unsafe: async (query: string, parameters?: string[]) => {
+      statements.push({ query, parameters })
+      if (query.includes('pg_extension')) return state.installed ? [{ '?column?': 1 }] : []
+      if (query.includes('pg_class'))
+        return (state.relations ?? []).map((relation) => ({ relation }))
+      if (query.includes('pg_prewarm(')) {
+        const [relation] = parameters ?? []
+        if (state.failing?.includes(relation))
+          throw new Error(`relation "${relation}" does not exist`)
+        return [{ pages: 7 }]
+      }
+      return []
+    },
+  }
+}
+
+describe('prewarmSearchProjection', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('does nothing where the extension is absent, so the application role never needs it', async () => {
+    const fake = session({ installed: false })
+    await expect(prewarmSearchProjection(fake)).resolves.toEqual([])
+    expect(fake.statements).toHaveLength(1)
+    expect(fake.statements[0].query).toContain("extname = 'pg_prewarm'")
+  })
+
+  it('reads the projections and their ranking indexes in the order the catalog lists them', async () => {
+    const fake = session({
+      installed: true,
+      relations: [
+        'embedding_search',
+        'embedding_keyword_tin',
+        'embedding_search_512_cosine_hnsw_idx',
+      ],
+    })
+    const warmed = await prewarmSearchProjection(fake)
+    expect(warmed.map((item) => item.relation)).toEqual([
+      'embedding_search',
+      'embedding_keyword_tin',
+      'embedding_search_512_cosine_hnsw_idx',
+    ])
+    expect(warmed.every((item) => item.pages === 7)).toBe(true)
+    const listed = fake.statements.find((statement) => statement.query.includes('pg_class'))
+    expect(listed?.parameters).toEqual([
+      '{embedding_search,embedding_keyword_tin}',
+      '{hnsw,tin,gin}',
+    ])
+    expect(listed?.query).toContain("ORDER BY c.relkind = 'r' DESC")
+    const reads = fake.statements.filter((statement) => statement.query.includes('pg_prewarm('))
+    expect(reads.map((statement) => statement.parameters)).toEqual([
+      ['embedding_search'],
+      ['embedding_keyword_tin'],
+      ['embedding_search_512_cosine_hnsw_idx'],
+    ])
+    expect(reads.every((statement) => statement.query.includes("'read'"))).toBe(true)
+  })
+
+  it('skips a relation that fails to warm and carries on with the rest', async () => {
+    const fake = session({
+      installed: true,
+      relations: ['embedding_search', 'embedding_search_512_cosine_hnsw_idx'],
+      failing: ['embedding_search'],
+    })
+    const warmed = await prewarmSearchProjection(fake)
+    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) => {
+      if (query.includes('pg_extension')) return [{ '?column?': 1 }]
+      throw new Error('permission denied for table pg_class')
+    }
+    await expect(prewarmSearchProjection(fake)).resolves.toEqual([])
+  })
+})
+
+describe('prewarmRelation', () => {
+  it('reports the pages read for one relation', async () => {
+    const fake = session({ installed: true })
+    await expect(prewarmRelation(fake, 'embedding_search')).resolves.toMatchObject({
+      relation: 'embedding_search',
+      pages: 7,
+    })
+    expect(await pgPrewarmInstalled(fake)).toBe(true)
+  })
+})

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

@@ -0,0 +1,154 @@
+import { PROJECTION_SOURCE_ACL_TABLES } from '@sim/db/script-migrations/0021_embedding_search_connector'
+import { createLogger } from '@sim/logger'
+import { getErrorMessage } from '@sim/utils/errors'
+
+const logger = createLogger('SearchProjectionPrewarm')
+
+/**
+ * The access methods a ranking touches at random: the vector graphs, the Tin keyword index, and
+ * the GIN index the on-row permission test reads. The remaining b-trees serve hydration, which
+ * reads a handful of rows by key and is fast cold.
+ */
+const RANKING_ACCESS_METHODS = ['hnsw', 'tin', 'gin'] as const
+
+/** The one call the helper needs from a `postgres` connection or a reserved session. */
+export interface PrewarmSession {
+  unsafe(query: string, parameters?: string[]): PromiseLike<ArrayLike<Record<string, unknown>>>
+}
+
+export interface PrewarmedRelation {
+  relation: string
+  pages: number
+  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
+ * touch it, which is what a bulk operation leaves behind.
+ */
+export async function pgPrewarmInstalled(session: PrewarmSession): Promise<boolean> {
+  const rows = await session.unsafe("SELECT 1 FROM pg_extension WHERE extname = 'pg_prewarm'")
+  return rows.length > 0
+}
+
+/**
+ * Reads one relation into the operating system's cache. `read` mode leaves shared buffers to the
+ * workload, where `buffer` mode would evict them wholesale to make room.
+ */
+export async function prewarmRelation(
+  session: PrewarmSession,
+  relation: string
+): Promise<PrewarmedRelation> {
+  const startedAt = Date.now()
+  const [row] = Array.from(
+    await session.unsafe("SELECT pg_prewarm($1::regclass, 'read')::int AS pages", [relation])
+  )
+  return { relation, pages: Number(row?.pages ?? 0), elapsedMs: Date.now() - startedAt }
+}
+
+/**
+ * Warms the ranking projections after something streamed through them. A backfill or index build
+ * reads every heap page in order and pushes the vector graphs out of cache; the next searches
+ * then fetch the graph one random page at a time from disk, take seconds, and end at their
+ * deadline with partial results. Reading the projections back in makes the first search after a
+ * bulk operation as fast as the thousandth.
+ *
+ * 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.
+ * 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,
+  options: PrewarmOptions = {}
+): Promise<PrewarmedRelation[]> {
+  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', {
+      error: getErrorMessage(error),
+    })
+    return []
+  }
+  const warmed: PrewarmedRelation[] = []
+  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: Date.now() - startedAt,
+  })
+  return warmed
+}
+
+/** The projections' heaps, then their ranking indexes smallest first, as the search path finds them. */
+async function rankingRelations(session: PrewarmSession): Promise<string[]> {
+  const rows = await session.unsafe(
+    `WITH heaps AS (
+       SELECT to_regclass(name) AS oid FROM unnest($1::text[]) AS name
+     )
+     SELECT c.oid::regclass::text AS relation
+     FROM pg_class c
+     JOIN pg_am am ON am.oid = c.relam
+     LEFT JOIN pg_index i ON i.indexrelid = c.oid
+     WHERE c.oid IN (SELECT oid FROM heaps)
+        OR (
+          i.indrelid IN (SELECT oid FROM heaps)
+          AND i.indisvalid
+          AND am.amname = ANY($2::text[])
+        )
+     ORDER BY c.relkind = 'r' DESC, pg_relation_size(c.oid)`,
+    [toArrayLiteral(PROJECTION_SOURCE_ACL_TABLES), toArrayLiteral(RANKING_ACCESS_METHODS)]
+  )
+  return Array.from(rows, (row) => String(row.relation))
+}
+
+/** Postgres array literal for identifiers that carry no quotes, commas or braces. */
+function toArrayLiteral(values: readonly string[]): string {
+  return `{${values.join(',')}}`
+}

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

@@ -3,10 +3,11 @@
  */
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockBackfill, mockEnd, mockPostgres, mockTasksTrigger } = vi.hoisted(() => ({
+const { mockBackfill, mockEnd, mockPostgres, mockPrewarm, mockTasksTrigger } = vi.hoisted(() => ({
   mockBackfill: vi.fn(),
   mockEnd: vi.fn(async () => undefined),
   mockPostgres: vi.fn(),
+  mockPrewarm: vi.fn(async () => []),
   mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })),
 }))
 
@@ -16,6 +17,7 @@ vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({
   backfillProjectionSourceAcl: mockBackfill,
 }))
 vi.mock('postgres', () => ({ default: mockPostgres }))
+vi.mock('@/lib/knowledge/search/prewarm', () => ({ prewarmSearchProjection: mockPrewarm }))
 vi.mock('@trigger.dev/sdk', () => ({ tasks: { trigger: mockTasksTrigger } }))
 vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' }))
 vi.mock('@/lib/core/utils/background', () => ({
@@ -26,6 +28,7 @@ vi.mock('@/lib/core/utils/background', () => ({
 
 import {
   enqueueProjectionSourceAclBackfill,
+  PROJECTION_PREWARM_BUDGET_MS,
   runProjectionSourceAclBackfill,
 } from '@/lib/knowledge/search/projection-source-acl-backfill'
 
@@ -57,6 +60,15 @@ describe('runProjectionSourceAclBackfill', () => {
     expect(mockEnd).toHaveBeenCalledTimes(1)
   })
 
+  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, { budgetMs: PROJECTION_PREWARM_BUDGET_MS })
+    expect(mockPrewarm.mock.invocationCallOrder[0]).toBeLessThan(
+      mockEnd.mock.invocationCallOrder[0]
+    )
+  })
+
   it('resumes after the cursor in its projection and from the start of the next', async () => {
     await runProjectionSourceAclBackfill({
       cursor: { projection: 'embedding_keyword_tin', afterId: 'chunk-9' },
@@ -80,6 +92,7 @@ describe('runProjectionSourceAclBackfill', () => {
     })
     expect(mockBackfill).toHaveBeenCalledTimes(1)
     expect(mockBackfill.mock.calls[0][2].budgetMs).toBeLessThanOrEqual(1000)
+    expect(mockPrewarm).not.toHaveBeenCalled()
     expect(mockEnd).toHaveBeenCalledTimes(1)
   })