|
| 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 | +/** |
| 26 | + * `pg_prewarm` is not a trusted extension, so the application role cannot create it and no |
| 27 | + * migration can; a superuser installs it once. Without it the projection warms only as searches |
| 28 | + * touch it, which is what a bulk operation leaves behind. |
| 29 | + */ |
| 30 | +export async function pgPrewarmInstalled(session: PrewarmSession): Promise<boolean> { |
| 31 | + const rows = await session.unsafe("SELECT 1 FROM pg_extension WHERE extname = 'pg_prewarm'") |
| 32 | + return rows.length > 0 |
| 33 | +} |
| 34 | + |
| 35 | +/** |
| 36 | + * Reads one relation into the operating system's cache. `read` mode leaves shared buffers to the |
| 37 | + * workload, where `buffer` mode would evict them wholesale to make room. |
| 38 | + */ |
| 39 | +export async function prewarmRelation( |
| 40 | + session: PrewarmSession, |
| 41 | + relation: string |
| 42 | +): Promise<PrewarmedRelation> { |
| 43 | + const startedAt = Date.now() |
| 44 | + const [row] = Array.from( |
| 45 | + await session.unsafe("SELECT pg_prewarm($1::regclass, 'read')::int AS pages", [relation]) |
| 46 | + ) |
| 47 | + return { relation, pages: Number(row?.pages ?? 0), elapsedMs: Date.now() - startedAt } |
| 48 | +} |
| 49 | + |
| 50 | +/** |
| 51 | + * Warms the ranking projections after something streamed through them. A backfill or index build |
| 52 | + * reads every heap page in order and pushes the vector graphs out of cache; the next searches |
| 53 | + * then fetch the graph one random page at a time from disk, take seconds, and end at their |
| 54 | + * deadline with partial results. Reading the projections back in makes the first search after a |
| 55 | + * bulk operation as fast as the thousandth. |
| 56 | + * |
| 57 | + * Heaps go first and the ranking indexes last, so where the cache cannot hold everything the |
| 58 | + * indexes are what survives: a walk reads far more index pages than heap pages. Relations are |
| 59 | + * resolved through the search path, so a schema that carries its own copy warms its own copy. |
| 60 | + * A relation that fails to warm is logged and skipped; warming is never worth failing the |
| 61 | + * operation that asked for it. |
| 62 | + */ |
| 63 | +export async function prewarmSearchProjection( |
| 64 | + session: PrewarmSession |
| 65 | +): Promise<PrewarmedRelation[]> { |
| 66 | + if (!(await pgPrewarmInstalled(session))) { |
| 67 | + logger.warn('pg_prewarm is not installed; the search projection warms only as it is searched') |
| 68 | + return [] |
| 69 | + } |
| 70 | + let relations: string[] |
| 71 | + try { |
| 72 | + relations = await rankingRelations(session) |
| 73 | + } catch (error) { |
| 74 | + logger.warn('Search projection relations could not be listed', { |
| 75 | + error: getErrorMessage(error), |
| 76 | + }) |
| 77 | + return [] |
| 78 | + } |
| 79 | + const warmed: PrewarmedRelation[] = [] |
| 80 | + for (const relation of relations) { |
| 81 | + try { |
| 82 | + warmed.push(await prewarmRelation(session, relation)) |
| 83 | + } catch (error) { |
| 84 | + logger.warn('Search projection relation failed to warm', { |
| 85 | + relation, |
| 86 | + error: getErrorMessage(error), |
| 87 | + }) |
| 88 | + } |
| 89 | + } |
| 90 | + logger.info('Search projection warmed', { |
| 91 | + relations: warmed.length, |
| 92 | + pages: warmed.reduce((sum, item) => sum + item.pages, 0), |
| 93 | + elapsedMs: warmed.reduce((sum, item) => sum + item.elapsedMs, 0), |
| 94 | + }) |
| 95 | + return warmed |
| 96 | +} |
| 97 | + |
| 98 | +/** The projections' heaps, then their ranking indexes smallest first, as the search path finds them. */ |
| 99 | +async function rankingRelations(session: PrewarmSession): Promise<string[]> { |
| 100 | + const rows = await session.unsafe( |
| 101 | + `WITH heaps AS ( |
| 102 | + SELECT to_regclass(name) AS oid FROM unnest($1::text[]) AS name |
| 103 | + ) |
| 104 | + SELECT c.oid::regclass::text AS relation |
| 105 | + FROM pg_class c |
| 106 | + JOIN pg_am am ON am.oid = c.relam |
| 107 | + LEFT JOIN pg_index i ON i.indexrelid = c.oid |
| 108 | + WHERE c.oid IN (SELECT oid FROM heaps) |
| 109 | + OR ( |
| 110 | + i.indrelid IN (SELECT oid FROM heaps) |
| 111 | + AND i.indisvalid |
| 112 | + AND am.amname = ANY($2::text[]) |
| 113 | + ) |
| 114 | + ORDER BY c.relkind = 'r' DESC, pg_relation_size(c.oid)`, |
| 115 | + [toArrayLiteral(PROJECTION_SOURCE_ACL_TABLES), toArrayLiteral(RANKING_ACCESS_METHODS)] |
| 116 | + ) |
| 117 | + return Array.from(rows, (row) => String(row.relation)) |
| 118 | +} |
| 119 | + |
| 120 | +/** Postgres array literal for identifiers that carry no quotes, commas or braces. */ |
| 121 | +function toArrayLiteral(values: readonly string[]): string { |
| 122 | + return `{${values.join(',')}}` |
| 123 | +} |
0 commit comments