mirror of
https://github.com/fluxerapp/fluxer.git
synced 2026-09-02 21:04:06 +03:00
fix(kv): restore keyset paging for numeric key scans (#2314)
This commit is contained in:
Hampus
@@ -68,6 +68,7 @@ class PlainClient implements IPostgresClient {
|
||||
|
||||
class CountingClient implements IPostgresClient {
|
||||
rowsRead = 0;
|
||||
readonly selects: Array<{text: string; values: Array<unknown>}> = [];
|
||||
|
||||
constructor(
|
||||
private readonly inner: IPostgresClient,
|
||||
@@ -75,7 +76,10 @@ class CountingClient implements IPostgresClient {
|
||||
) {}
|
||||
async query<T extends Record<string, unknown>>(text: string, values: Array<unknown> = [], 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<ReturnType<IPostgresClient['query']>> & {rows: Array<T>};
|
||||
}
|
||||
async connect(): Promise<void> {
|
||||
@@ -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<number> = [];
|
||||
let pageState: string | null = null;
|
||||
let seen = 0;
|
||||
for (let guard = 0; guard < 100; guard += 1) {
|
||||
const before = counter.rowsRead;
|
||||
const page: {rows: Array<Row>; pageState: string | null} = await exec.executePagedQuery<Row>(
|
||||
{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<number>) => 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<Row>(query, {pageSize: 4});
|
||||
expect(first.pageState).not.toBeNull();
|
||||
counter.selects.length = 0;
|
||||
await exec.executePagedQuery<Row>(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<Record<string, string>>(`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<bigint>) => {
|
||||
await wipe(KV);
|
||||
|
||||
@@ -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<string>).length === 1;
|
||||
}
|
||||
|
||||
function numericScanKeyed(meta: KvQueryMeta, plan: QueryPlan, entries: ReadonlyArray<PageEntry>): boolean {
|
||||
if (!numericScanPlan(meta, plan)) return false;
|
||||
const column = (meta.table.primaryKey as ReadonlyArray<string>)[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<PageEntry>, 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<T>, pageState: page.pageState};
|
||||
}
|
||||
const rows = await this.executeQuery<T, P>(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<Row>; pageState: string | null}> {
|
||||
logFullScan(meta);
|
||||
const result = await this.client.query<StoredRow>(
|
||||
`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<string> | 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<PageEntry>): Array<Row> {
|
||||
return entries.map((entry) => projectRow(entry.row, meta.columns as ReadonlyArray<string> | undefined));
|
||||
}
|
||||
|
||||
async executeBatch(
|
||||
queries: Array<{query: string; params: object; meta?: KvQueryMeta}>,
|
||||
atomic = true,
|
||||
|
||||
@@ -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]);
|
||||
|
||||
Reference in New Issue
Block a user