From 42df4f673144f94f144375ad7ee7cebc82345d33 Mon Sep 17 00:00:00 2001 From: Hampus Date: Tue, 1 Sep 2026 01:32:33 +0200 Subject: [PATCH] fix(kv): drop the row_key order probe that cliffed paged scans (#2306) --- fluxer_api/src/api/app/APILifecycle.ts | 7 +- .../src/api/app/DeletionQueueStartup.ts | 30 ++--- .../app/tests/DeletionQueueStartup.test.ts | 17 +++ .../PostgresKvPagingAdversarial.test.ts | 74 +++++++----- .../api/database/PostgresKvQueryExecutor.ts | 107 +----------------- 5 files changed, 82 insertions(+), 153 deletions(-) diff --git a/fluxer_api/src/api/app/APILifecycle.ts b/fluxer_api/src/api/app/APILifecycle.ts index 9baa77623..3707caa81 100644 --- a/fluxer_api/src/api/app/APILifecycle.ts +++ b/fluxer_api/src/api/app/APILifecycle.ts @@ -134,12 +134,7 @@ export function createInitializer(config: APIConfig, logger: ILogger): () => Pro setInjectedWorkerService(new WorkerService(workerQueue, getSnowflakeService(), new JobLedgerRepository())); logger.info('JetStream worker service initialized'); } - try { - await ensureDeletionQueueState(getKVAccountDeletionQueue(), logger); - } catch (error) { - logger.error({error}, 'Failed to verify KV deletion queue state'); - throw error; - } + await ensureDeletionQueueState(getKVAccountDeletionQueue(), logger); logger.info('Initializing search indexes...'); let searchInitialized = false; try { diff --git a/fluxer_api/src/api/app/DeletionQueueStartup.ts b/fluxer_api/src/api/app/DeletionQueueStartup.ts index ca313207d..8c7eeb3e1 100644 --- a/fluxer_api/src/api/app/DeletionQueueStartup.ts +++ b/fluxer_api/src/api/app/DeletionQueueStartup.ts @@ -7,19 +7,23 @@ export async function ensureDeletionQueueState( deletionQueue: KVAccountDeletionQueueService, logger: ILogger, ): Promise { - if (!(await deletionQueue.needsRebuild())) { - logger.info('KV deletion queue state is healthy'); - return; - } - const lockToken = await deletionQueue.acquireRebuildLock(); - if (!lockToken) { - logger.info('Another instance is rebuilding the KV deletion queue, skipping'); - return; - } - logger.info('KV deletion queue needs rebuild, rebuilding...'); try { - await deletionQueue.rebuildState(lockToken); - } finally { - await deletionQueue.releaseRebuildLock(lockToken); + if (!(await deletionQueue.needsRebuild())) { + logger.info('KV deletion queue state is healthy'); + return; + } + const lockToken = await deletionQueue.acquireRebuildLock(); + if (!lockToken) { + logger.info('Another instance is rebuilding the KV deletion queue, skipping'); + return; + } + logger.info('KV deletion queue needs rebuild, rebuilding...'); + try { + await deletionQueue.rebuildState(lockToken); + } finally { + await deletionQueue.releaseRebuildLock(lockToken); + } + } catch (error) { + logger.error({error}, 'KV deletion queue rebuild failed, continuing startup'); } } diff --git a/fluxer_api/src/api/app/tests/DeletionQueueStartup.test.ts b/fluxer_api/src/api/app/tests/DeletionQueueStartup.test.ts index b3be180ca..3d19bceb6 100644 --- a/fluxer_api/src/api/app/tests/DeletionQueueStartup.test.ts +++ b/fluxer_api/src/api/app/tests/DeletionQueueStartup.test.ts @@ -72,6 +72,23 @@ describe('ensureDeletionQueueState', () => { expect(apiFailures).toEqual([]); }); + it('does not abort startup when the paged user scan fails', async () => { + const kvClient = new MockKVProvider(); + let scans = 0; + const repository = { + async scanAllUsersPage() { + scans += 1; + throw new Error('paged user scan failed'); + }, + } as unknown as UserRepository; + const queue = new KVAccountDeletionQueueService(kvClient, repository); + + await expect(ensureDeletionQueueState(queue, new NoopLogger())).resolves.toBeUndefined(); + + expect(scans).toBe(1); + expect(await queue.acquireRebuildLock()).not.toBeNull(); + }); + it('rebuilds under the lock when no other instance holds it', async () => { const kvClient = new MockKVProvider(); const queue = new KVAccountDeletionQueueService( diff --git a/fluxer_api/src/api/database/PostgresKvPagingAdversarial.test.ts b/fluxer_api/src/api/database/PostgresKvPagingAdversarial.test.ts index e2921a06c..ef0928de8 100644 --- a/fluxer_api/src/api/database/PostgresKvPagingAdversarial.test.ts +++ b/fluxer_api/src/api/database/PostgresKvPagingAdversarial.test.ts @@ -574,46 +574,60 @@ 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: `key${String(i).padStart(2, '0')}`, v: `v${i}`}); - } - const drain = async (build: (client: IPostgresClient) => AnyExec) => { + 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); + for (const id of ids) await upsert(next, FlatTable, {k: id, v: id.toString()}); const counter = new CountingClient(raw, KV); - const exec = build(counter); - const reads: Array = []; + const exec = new PostgresKvQueryExecutor(counter); + const seen: 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}, + {cql: '__digits__', params: {}, kvMeta: selectMeta(FlatTable)}, + {pageSize: 4, pageState}, ); - reads.push(counter.rowsRead - before); - seen += page.rows.length; + for (const row of page.rows) seen.push(String(row.k)); pageState = page.pageState; if (pageState === null) break; } - return {reads, seen}; + return {seen, rowsRead: counter.rowsRead}; }; - 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); + const sameWidth = Array.from({length: 40}, (_, index) => 1_000_000_000_000_000_000n + BigInt(index)); + const straddling = Array.from({length: 40}, (_, index) => 999_999_999_999_999_980n + BigInt(index)); + const flat = await drain(sameWidth); + const crossing = await drain(straddling); + expect(flat.seen, 'same width scan order').toEqual(sameWidth.map(String)); + expect(crossing.seen, 'boundary crossing scan order').toEqual(straddling.map(String)); 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, - ); + crossing.rowsRead, + `same width read ${flat.rowsRead} rows, boundary crossing read ${crossing.rowsRead}`, + ).toBe(flat.rowsRead); + }, 300_000); + + it('keeps paging bigint keys across a digit count boundary while returned rows are deleted', async () => { + await wipe(KV); + const ids = Array.from({length: 30}, (_, index) => 999_999_999_999_999_985n + BigInt(index)); + for (const id of ids) await upsert(next, FlatTable, {k: id, v: id.toString()}); + const seen: Array = []; + let pageState: string | null = null; + for (let guard = 0; guard < 100; guard += 1) { + const page: {rows: Array; pageState: string | null} = await next.executePagedQuery( + {cql: '__digitdrain__', params: {}, kvMeta: selectMeta(FlatTable)}, + {pageSize: 4, pageState}, + ); + for (const row of page.rows) seen.push(String(row.k)); + pageState = page.pageState; + for (const row of page.rows) { + await next.executeQuery({ + cql: '__digitdrain_delete__', + params: {k: row.k} as CassandraParams, + kvMeta: deleteMeta(FlatTable, [{kind: 'eq', col: 'k', param: 'k'} as WhereExpr]), + }); + } + if (pageState === null) break; + } + expect(seen, `returned order: ${seen.join(',')}`).toEqual(ids.map(String)); }, 300_000); it('pages a prefix range whose keys sit adjacent to the range bounds', async () => { diff --git a/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts b/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts index 1e5aacd66..9d166714a 100644 --- a/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts +++ b/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts @@ -19,7 +19,6 @@ interface StoredRow { interface PageState { offset: number; after?: string; - keyed?: boolean; } interface PageEntry { @@ -65,7 +64,6 @@ const VALUE_SEPARATOR = '\u001f'; const KEY_RANGE_UPPER = ' '; const MAX_ROW_KEY_COMBINATIONS = 32_768; const MAX_PREFIX_RANGES = 256; -const PAGE_PROBE_BATCH = 10_000; const FULL_SCAN_LOG_INTERVAL_MS = 60_000; const FULL_SCAN_LOG_KEY_LIMIT = 1024; const ENCODED_TYPE_KEY = '__fluxer_type'; @@ -671,7 +669,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, keyed: decoded.keyed === true}; + return {offset: decoded.offset, after: decoded.after}; } function pageableSelect(meta: KvQueryMeta, pageSize: number): boolean { @@ -684,13 +682,6 @@ function pageableSelect(meta: KvQueryMeta, pageSize: number): boolean { ); } -function keysetCandidates(plan: QueryPlan): boolean { - return ( - plan.exact && - (plan.candidates.kind === 'none' || plan.candidates.kind === 'range' || plan.candidates.kind === 'scan') - ); -} - 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); @@ -702,12 +693,6 @@ function pageStart(meta: KvQueryMeta, entries: ReadonlyArray, state: return start; } -function pagedStatementName(prefix: string, plan: CandidatePlan, after: string | null): string | undefined { - const name = planStatementName(prefix, plan); - if (name === undefined) return undefined; - return after === null ? name : `${name}_after`; -} - function parseRawMeta(cql: string): KvQueryMeta | null { const normalized = normalizeCql(cql).replace(/;$/, ''); const update = @@ -930,7 +915,8 @@ export class PostgresKvQueryExecutor { const state = decodePageState(options.pageState); const meta = this.meta(query); if (pageableSelect(meta, options.pageSize) && (state.after !== undefined || state.offset === 0)) { - const page = await this.selectPage(meta, query.params, state, options.pageSize); + const plan = buildCandidatePlan(meta, query.params); + const page = 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); @@ -942,93 +928,6 @@ export class PostgresKvQueryExecutor { }; } - private async selectPage( - meta: KvQueryMeta, - params: CassandraParams, - state: PageState, - pageSize: number, - ): Promise<{rows: Array; pageState: string | null}> { - const plan = buildCandidatePlan(meta, params); - if (keysetCandidates(plan)) { - if (state.keyed === true) { - return this.keysetPage(meta, params, plan, state.after ?? null, state.offset, pageSize); - } - if (state.after === undefined && (await this.rowKeyOrderMatchesSort(meta, plan))) { - return this.keysetPage(meta, params, plan, null, 0, pageSize); - } - } - return this.sortedPage(meta, params, plan, state, pageSize); - } - - private pagedSql( - meta: KvQueryMeta, - fragments: PlanFragments, - projection: string, - after: string | null, - limit: number, - ): {text: string; values: Array} { - const values: Array = [meta.table.name, ...fragments.params]; - let text = `SELECT ${projection} FROM ${this.table} kv WHERE kv.table_name = $1${fragments.predicate}`; - if (after !== null) { - values.push(after); - text += ` AND kv.row_key COLLATE "C" > $${values.length}`; - } - text += ' AND (kv.expires_at IS NULL OR kv.expires_at > now()) ORDER BY kv.row_key COLLATE "C"'; - values.push(limit); - return {text: `${text} LIMIT $${values.length}`, values}; - } - - private async rowKeyOrderMatchesSort(meta: KvQueryMeta, plan: QueryPlan): Promise { - if (plan.candidates.kind === 'none') return true; - const fragments = planFragmentGroups(plan.candidates)[0]!; - const columns = (meta.table.primaryKey as ReadonlyArray).length; - let after: string | null = null; - let previous: Array | null = null; - for (;;) { - const sql = this.pagedSql(meta, fragments, 'kv.row_key', after, PAGE_PROBE_BATCH); - const name = pagedStatementName('kv_keys', plan.candidates, after); - const batch: Array<{row_key: string}> = (await this.client.query<{row_key: string}>(sql.text, sql.values, name)) - .rows; - for (const stored of batch) { - const values = decodeRowKey(stored.row_key, columns); - if (values === null) return false; - if (previous !== null && compareKeyValues(previous, values) > 0) return false; - previous = values; - } - if (batch.length < PAGE_PROBE_BATCH) return true; - after = batch[batch.length - 1]!.row_key; - } - } - - private async keysetPage( - meta: KvQueryMeta, - params: CassandraParams, - plan: QueryPlan, - after: string | null, - offset: number, - pageSize: number, - ): Promise<{rows: Array; pageState: string | null}> { - if (plan.candidates.kind === 'none') return {rows: [], pageState: null}; - if (plan.candidates.kind === 'scan') logFullScan(meta); - const fragments = planFragmentGroups(plan.candidates)[0]!; - const sql = this.pagedSql(meta, fragments, 'kv.row_key, kv.row_data', after, pageSize + 1); - const result = await this.client.query( - sql.text, - sql.values, - pagedStatementName('kv_page', plan.candidates, after), - ); - const stored = result.rows.slice(0, pageSize); - const entries = this.matchingEntries(meta, stored, params); - const last = entries[entries.length - 1]; - return { - rows: entries.map((entry) => projectRow(entry.row, meta.columns as ReadonlyArray | undefined)), - pageState: - result.rows.length > pageSize && last - ? encodePageState({offset: offset + entries.length, after: last.key, keyed: true}) - : null, - }; - } - private async sortedPage( meta: KvQueryMeta, params: CassandraParams,