From 24cd163acdc1018f49103cc838f563f1c82b88ba Mon Sep 17 00:00:00 2001 From: Hampus Date: Tue, 1 Sep 2026 02:54:10 +0200 Subject: [PATCH] fix(kv): restore keyset paging for numeric key scans (#2314) --- .../PostgresKvPagingAdversarial.test.ts | 69 +++++++++++++++- .../api/database/PostgresKvQueryExecutor.ts | 78 ++++++++++++++++++- .../database/PostgresKvUpgradeSafety.test.ts | 7 +- 3 files changed, 148 insertions(+), 6 deletions(-) diff --git a/fluxer_api/src/api/database/PostgresKvPagingAdversarial.test.ts b/fluxer_api/src/api/database/PostgresKvPagingAdversarial.test.ts index ef0928de8..791215dfc 100644 --- a/fluxer_api/src/api/database/PostgresKvPagingAdversarial.test.ts +++ b/fluxer_api/src/api/database/PostgresKvPagingAdversarial.test.ts @@ -68,6 +68,7 @@ class PlainClient implements IPostgresClient { class CountingClient implements IPostgresClient { rowsRead = 0; + readonly selects: Array<{text: string; values: Array}> = []; constructor( private readonly inner: IPostgresClient, @@ -75,7 +76,10 @@ class CountingClient implements IPostgresClient { ) {} async query>(text: string, values: Array = [], name?: string) { const result = await this.inner.query(text, values, name); - if (/^\s*SELECT\s+(kv\.)?row_key/iu.test(text) && text.includes(this.table)) this.rowsRead += result.rows.length; + if (/^\s*SELECT\s+(kv\.)?row_key/iu.test(text) && text.includes(this.table)) { + this.rowsRead += result.rows.length; + this.selects.push({text, values}); + } return result as unknown as Awaited> & {rows: Array}; } async connect(): Promise { @@ -574,6 +578,69 @@ describe.skipIf(!dockerAvailable)('postgres kv paging adversarial', () => { expect(deltas, `paged order deltas:\n${deltas.join('\n')}`).toEqual([]); }, 300_000); + it('reads a page instead of the whole table for every page after the first', async () => { + await wipe(KV); + const rowCount = 40; + const pageSize = 4; + for (let i = 0; i < rowCount; i += 1) { + await upsert(next, FlatTable, {k: 999_999_999_999_999_980n + BigInt(i), v: `v${i}`}); + } + const drain = async (build: (client: IPostgresClient) => AnyExec) => { + const counter = new CountingClient(raw, KV); + const exec = build(counter); + const reads: Array = []; + let pageState: string | null = null; + let seen = 0; + for (let guard = 0; guard < 100; guard += 1) { + const before = counter.rowsRead; + const page: {rows: Array; pageState: string | null} = await exec.executePagedQuery( + {cql: '__reads__', params: {}, kvMeta: selectMeta(FlatTable)}, + {pageSize, pageState}, + ); + reads.push(counter.rowsRead - before); + seen += page.rows.length; + pageState = page.pageState; + if (pageState === null) break; + } + return {reads, seen}; + }; + const legacyDrain = await drain((client) => new LegacyPostgresKvQueryExecutor(client)); + const nextDrain = await drain((client) => new PostgresKvQueryExecutor(client)); + const total = (reads: Array) => reads.reduce((sum, count) => sum + count, 0); + expect(legacyDrain.seen).toBe(rowCount); + expect(nextDrain.seen).toBe(rowCount); + expect(nextDrain.reads.length).toBe(rowCount / pageSize); + expect(Math.min(...legacyDrain.reads), 'offset paging re-reads the whole table for every page').toBe(rowCount); + expect( + Math.max(...nextDrain.reads.slice(1)), + `rows read per page: ${nextDrain.reads.join(',')}`, + ).toBeLessThanOrEqual(pageSize + 1); + expect(total(nextDrain.reads), `next=${total(nextDrain.reads)} legacy=${total(legacyDrain.reads)}`).toBeLessThan( + total(legacyDrain.reads) / 2, + ); + }, 300_000); + + it('serves every page after the first from an index instead of sorting the table', async () => { + await wipe(KV); + for (let i = 0; i < 2000; i += 1) { + await upsert(next, FlatTable, {k: 999_999_999_999_999_000n + BigInt(i), v: `v${i}`}); + } + await raw.query(`ANALYZE ${KV}`); + const counter = new CountingClient(raw, KV); + const exec = new PostgresKvQueryExecutor(counter); + const query = {cql: '__plan__', params: {}, kvMeta: selectMeta(FlatTable)}; + const first = await exec.executePagedQuery(query, {pageSize: 4}); + expect(first.pageState).not.toBeNull(); + counter.selects.length = 0; + await exec.executePagedQuery(query, {pageSize: 4, pageState: first.pageState}); + const paged = counter.selects[counter.selects.length - 1]; + expect(paged, 'no paged select was issued').toBeDefined(); + const explained = await raw.query>(`EXPLAIN ${paged!.text}`, paged!.values); + const plan = explained.rows.map((row) => Object.values(row)[0]).join('\n'); + expect(plan, plan).toContain(`${KV}_row_key_numeric_idx`); + expect(plan, plan).not.toContain('Seq Scan'); + }, 300_000); + it('pages a bigint scan at the same cost on both sides of a digit count boundary', async () => { const drain = async (ids: ReadonlyArray) => { await wipe(KV); diff --git a/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts b/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts index 9d166714a..ffb487c41 100644 --- a/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts +++ b/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts @@ -19,6 +19,7 @@ interface StoredRow { interface PageState { offset: number; after?: string; + keyed?: boolean; } interface PageEntry { @@ -72,10 +73,18 @@ const POSTGRES_KV_SCHEMA_LOCK_TIMEOUT = '120s'; const POSTGRES_KV_SCHEMA_MIGRATION_TIMEOUT = '30min'; export const POSTGRES_KV_MIGRATION_TABLE = '__fluxer_schema_migrations'; const POSTGRES_KV_MESSAGES_PARTITION_MIGRATION = 'messages_partition_key_v1'; +const NUMERIC_ROW_KEY_BIGINT_PATTERN = '^\\{"__fluxer_type":"bigint","value":"(-?[0-9]+)"\\}$'; +const NUMERIC_ROW_KEY_NUMBER_PATTERN = '^(-?[0-9]+(?:\\.[0-9]+)?(?:[eE][-+]?[0-9]+)?)$'; const EXPIRED_STORED_ROW = 'kv.expires_at IS NOT NULL AND kv.expires_at <= now()'; const MERGED_ROW_DATA = `CASE WHEN ${EXPIRED_STORED_ROW} THEN EXCLUDED.row_data ELSE kv.row_data || EXCLUDED.row_data END`; const KEPT_EXPIRES_AT = `CASE WHEN ${EXPIRED_STORED_ROW} THEN NULL ELSE kv.expires_at END`; +function numericRowKeyExpr(column: string): string { + return `(COALESCE(substring(${column} from '${NUMERIC_ROW_KEY_BIGINT_PATTERN}'), substring(${column} from '${NUMERIC_ROW_KEY_NUMBER_PATTERN}'))::numeric)`; +} + +const NUMERIC_ROW_KEY = numericRowKeyExpr('kv.row_key'); + function planStatementName(prefix: string, plan: CandidatePlan): string | undefined { switch (plan.kind) { case 'rowKeys': @@ -669,7 +678,7 @@ function decodePageState(pageState: string | null | undefined): PageState { throw new Error('Invalid Postgres KV page state'); } if (typeof decoded.after !== 'string') return {offset: decoded.offset}; - return {offset: decoded.offset, after: decoded.after}; + return {offset: decoded.offset, after: decoded.after, keyed: decoded.keyed === true}; } function pageableSelect(meta: KvQueryMeta, pageSize: number): boolean { @@ -682,6 +691,30 @@ function pageableSelect(meta: KvQueryMeta, pageSize: number): boolean { ); } +function numericKeyValue(value: unknown): string | null { + if (typeof value === 'bigint') return value.toString(); + if (typeof value === 'number' && Number.isFinite(value)) return String(value); + return null; +} + +function numericScanPlan(meta: KvQueryMeta, plan: QueryPlan): boolean { + return plan.exact && plan.candidates.kind === 'scan' && (meta.table.primaryKey as ReadonlyArray).length === 1; +} + +function numericScanKeyed(meta: KvQueryMeta, plan: QueryPlan, entries: ReadonlyArray): boolean { + if (!numericScanPlan(meta, plan)) return false; + const column = (meta.table.primaryKey as ReadonlyArray)[0]!; + return entries.every((entry) => numericKeyValue(entry.row[column]) !== null); +} + +function numericScanCursor(state: PageState): {rowKey: string; value: string} | null { + if (state.keyed !== true || state.after === undefined) return null; + const values = decodeRowKey(state.after, 1); + if (values === null) return null; + const value = numericKeyValue(values[0]); + return value === null ? null : {rowKey: state.after, value}; +} + function pageStart(meta: KvQueryMeta, entries: ReadonlyArray, state: PageState): number { if (state.after === undefined) return state.offset; const index = entries.findIndex((entry) => entry.key === state.after); @@ -808,6 +841,9 @@ CREATE TABLE IF NOT EXISTS ${table} ( `CREATE INDEX IF NOT EXISTS ${quoteIdentifier(`${kvTable}_row_key_c_idx`)} ON ${table} (table_name, row_key COLLATE "C")`, ); } + await db.query( + `CREATE INDEX IF NOT EXISTS ${quoteIdentifier(`${kvTable}_row_key_numeric_idx`)} ON ${table} (table_name, ${numericRowKeyExpr('row_key')}) WHERE ${numericRowKeyExpr('row_key')} IS NOT NULL`, + ); await db.query( `CREATE INDEX IF NOT EXISTS ${quoteIdentifier(`${kvTable}_expires_idx`)} ON ${table} (expires_at) WHERE expires_at IS NOT NULL`, ); @@ -916,7 +952,10 @@ export class PostgresKvQueryExecutor { const meta = this.meta(query); if (pageableSelect(meta, options.pageSize) && (state.after !== undefined || state.offset === 0)) { const plan = buildCandidatePlan(meta, query.params); - const page = await this.sortedPage(meta, query.params, plan, state, options.pageSize); + const cursor = numericScanPlan(meta, plan) ? numericScanCursor(state) : null; + const page = cursor + ? await this.numericScanPage(meta, query.params, cursor, state.offset, options.pageSize) + : await this.sortedPage(meta, query.params, plan, state, options.pageSize); return {rows: page.rows as Array, pageState: page.pageState}; } const rows = await this.executeQuery(query); @@ -928,6 +967,29 @@ export class PostgresKvQueryExecutor { }; } + private async numericScanPage( + meta: KvQueryMeta, + params: CassandraParams, + cursor: {rowKey: string; value: string}, + offset: number, + pageSize: number, + ): Promise<{rows: Array; pageState: string | null}> { + logFullScan(meta); + const result = await this.client.query( + `SELECT kv.row_key, kv.row_data FROM ${this.table} kv WHERE kv.table_name = $1 AND ${NUMERIC_ROW_KEY} IS NOT NULL AND (${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C") > ($2::numeric, $3) AND (kv.expires_at IS NULL OR kv.expires_at > now()) ORDER BY ${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C" LIMIT $4`, + [meta.table.name, cursor.value, cursor.rowKey, pageSize + 1], + ); + const entries = this.matchingEntries(meta, result.rows.slice(0, pageSize), params); + const last = entries[entries.length - 1]; + return { + rows: this.projected(meta, entries), + pageState: + result.rows.length > pageSize && last + ? encodePageState({offset: offset + entries.length, after: last.key, keyed: true}) + : null, + }; + } + private async sortedPage( meta: KvQueryMeta, params: CassandraParams, @@ -942,12 +1004,20 @@ export class PostgresKvQueryExecutor { const page = entries.slice(start, start + pageSize); const last = page[page.length - 1]; const nextOffset = start + page.length; + if (nextOffset >= entries.length || !last) return {rows: this.projected(meta, page), pageState: null}; + const keyed = numericScanKeyed(meta, plan, entries); return { - rows: page.map((entry) => projectRow(entry.row, meta.columns as ReadonlyArray | undefined)), - pageState: nextOffset < entries.length && last ? encodePageState({offset: nextOffset, after: last.key}) : null, + rows: this.projected(meta, page), + pageState: encodePageState( + keyed ? {offset: nextOffset, after: last.key, keyed: true} : {offset: nextOffset, after: last.key}, + ), }; } + private projected(meta: KvQueryMeta, entries: ReadonlyArray): Array { + return entries.map((entry) => projectRow(entry.row, meta.columns as ReadonlyArray | undefined)); + } + async executeBatch( queries: Array<{query: string; params: object; meta?: KvQueryMeta}>, atomic = true, diff --git a/fluxer_api/src/api/database/PostgresKvUpgradeSafety.test.ts b/fluxer_api/src/api/database/PostgresKvUpgradeSafety.test.ts index dcec88fd9..08e29210e 100644 --- a/fluxer_api/src/api/database/PostgresKvUpgradeSafety.test.ts +++ b/fluxer_api/src/api/database/PostgresKvUpgradeSafety.test.ts @@ -306,9 +306,13 @@ suite('postgres kv upgrade safety', () => { WHERE tablename = '${OLD}' AND indexname <> '${OLD}_row_key_c_idx' EXCEPT SELECT replace(indexdef, '${NEW}', 'KV') FROM pg_indexes WHERE tablename = '${NEW}') UNION ALL - (SELECT replace(indexdef, '${NEW}', 'KV') FROM pg_indexes WHERE tablename = '${NEW}' + (SELECT replace(indexdef, '${NEW}', 'KV') FROM pg_indexes + WHERE tablename = '${NEW}' AND indexname <> '${NEW}_row_key_numeric_idx' EXCEPT SELECT replace(indexdef, '${OLD}', 'KV') FROM pg_indexes WHERE tablename = '${OLD}') ) d`); + const added = await raw.query<{indexname: string}>(` + SELECT indexname FROM pg_indexes WHERE tablename = '${NEW}' + EXCEPT SELECT replace(indexname, '${OLD}', '${NEW}') FROM pg_indexes WHERE tablename = '${OLD}'`); const collations = await raw.query<{tablename: string; attname: string; collname: string}>(` SELECT cls.relname AS tablename, att.attname, col.collname FROM pg_attribute att @@ -322,6 +326,7 @@ suite('postgres kv upgrade safety', () => { ); expect(Number(diff.rows[0]!.n)).toBe(0); expect(Number(schemaDiff.rows[0]!.n)).toBe(0); + expect(added.rows.map((r) => r.indexname)).toEqual([`${NEW}_row_key_numeric_idx`]); expect(collations.rows.filter((r) => r.tablename === OLD).map((r) => r.collname)).toEqual(['default', 'default']); expect(collations.rows.filter((r) => r.tablename === NEW).map((r) => r.collname)).toEqual(['C', 'C']); expect(cIndexes.rows.map((r) => r.tablename)).toEqual([OLD]);