fix(kv): page scans by keyset so deletes cannot skip rows (#2302)

fluxer_api/src/api/database/PostgresKvPagingAdversarial.test.ts

@@ -66,6 +66,33 @@ class PlainClient implements IPostgresClient {
 	}
 }
 
+class CountingClient implements IPostgresClient {
+	rowsRead = 0;
+
+	constructor(
+		private readonly inner: IPostgresClient,
+		private readonly table: string,
+	) {}
+	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;
+		return result as unknown as Awaited<ReturnType<IPostgresClient['query']>> & {rows: Array<T>};
+	}
+	async connect(): Promise<void> {
+		await this.inner.connect();
+	}
+	async shutdown(): Promise<void> {}
+	isConnected(): boolean {
+		return this.inner.isConnected();
+	}
+	async transaction<T>(fn: (client: PostgresQueryable) => Promise<T>): Promise<T> {
+		return this.inner.transaction(fn);
+	}
+	kvTable(): string {
+		return this.table;
+	}
+}
+
 function selectMeta(spec: KvTableSpec<Row>, where: Array<WhereExpr<Row>> = [], extra: Partial<AnyMeta> = {}): AnyMeta {
 	return {action: 'select', table: spec, where, columns: spec.columns, ...extra} as AnyMeta;
 }
@@ -247,7 +274,7 @@ describe.skipIf(!dockerAvailable)('postgres kv paging adversarial', () => {
 		expect(dupes, `duplicate rows: ${dupes.join(',')}`).toEqual([]);
 	}, 300_000);
 
-	it('pages identically to the legacy executor while rows are deleted between pages', async () => {
+	it('keeps returning later rows while already-returned rows are deleted between pages', async () => {
 		const run = async (exec: AnyExec) => {
 			await wipe(KV);
 			for (let i = 0; i < 30; i += 1) await upsert(exec, FlatTable, {k: BigInt(i), v: `v${i}`});
@@ -261,18 +288,23 @@ describe.skipIf(!dockerAvailable)('postgres kv paging adversarial', () => {
 					kvMeta: deleteMeta(FlatTable, [{kind: 'eq', col: 'k', param: 'k'} as WhereExpr<Row>]),
 				});
 			});
-			return {
-				error: result.error,
-				pages: result.pages.map((page) => page.map((row) => String(row.k)).join(',')).join('|'),
-			};
+			return {error: result.error, seen: result.pages.flat().map((row) => String(row.k))};
 		};
 		const legacyRun = await run(legacy);
 		const nextRun = await run(next);
+		const seeded = Array.from({length: 30}, (_, index) => String(index));
 		expect(nextRun.error).toBeNull();
-		expect(nextRun, `legacy=${legacyRun.pages} next=${nextRun.pages}`).toEqual(legacyRun);
+		const skipped = seeded.filter((key) => !nextRun.seen.includes(key));
+		const dupes = nextRun.seen.filter((key, index) => nextRun.seen.indexOf(key) !== index);
+		expect(skipped, `rows never returned: ${skipped.join(',')}`).toEqual([]);
+		expect(dupes, `rows returned twice: ${dupes.join(',')}`).toEqual([]);
+		expect(
+			seeded.filter((key) => !legacyRun.seen.includes(key)).length,
+			'offset paging is expected to skip rows once earlier rows are deleted',
+		).toBeGreaterThan(0);
 	}, 300_000);
 
-	it('pages identically to the legacy executor when the last row of each page is deleted', async () => {
+	it('does not skip the row after a page cursor that was deleted', async () => {
 		const run = async (exec: AnyExec) => {
 			await wipe(KV);
 			for (let i = 0; i < 20; i += 1)
@@ -299,9 +331,17 @@ describe.skipIf(!dockerAvailable)('postgres kv paging adversarial', () => {
 			}
 			return seen.join(',');
 		};
-		const legacySeen = await run(legacy);
-		const nextSeen = await run(next);
-		expect(nextSeen, `legacy=${legacySeen} next=${nextSeen}`).toBe(legacySeen);
+		const legacySeen = (await run(legacy)).split(',');
+		const nextSeen = (await run(next)).split(',');
+		const seeded = Array.from({length: 20}, (_, index) => `key${String(index).padStart(2, '0')}`);
+		const skipped = seeded.filter((key) => !nextSeen.includes(key));
+		const dupes = nextSeen.filter((key, index) => nextSeen.indexOf(key) !== index);
+		expect(skipped, `rows never returned: ${skipped.join(',')}`).toEqual([]);
+		expect(dupes, `rows returned twice: ${dupes.join(',')}`).toEqual([]);
+		expect(
+			seeded.filter((key) => !legacySeen.includes(key)).length,
+			'offset paging is expected to skip the row after a deleted cursor',
+		).toBeGreaterThan(0);
 	}, 300_000);
 
 	it('never returns an empty page together with a non-null page state', async () => {
@@ -534,6 +574,48 @@ 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) => {
+			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('pages a prefix range whose keys sit adjacent to the range bounds', async () => {
 		await wipe(KV);
 		const owners = ['a', 'a b', 'ab', 'a
', 'a"'];

fluxer_api/src/api/database/PostgresKvQueryExecutor.ts

@@ -18,6 +18,13 @@ interface StoredRow {
 
 interface PageState {
 	offset: number;
+	after?: string;
+	keyed?: boolean;
+}
+
+interface PageEntry {
+	key: string;
+	row: Row;
 }
 
 interface RangeGroup {
@@ -58,6 +65,7 @@ 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';
@@ -296,18 +304,53 @@ function projectRow(row: Row, columns: ReadonlyArray<string> | undefined): Row {
 	return projected;
 }
 
-function sortRows(meta: KvQueryMeta, rows: Array<Row>): Array<Row> {
+function rowComparator(meta: KvQueryMeta): (left: Row, right: Row) => number {
 	if (meta.orderBy) {
+		const column = meta.orderBy.col as string;
 		const direction = meta.orderBy.direction === 'DESC' ? -1 : 1;
-		return rows.sort((left, right) => compareValues(left[meta.orderBy!.col], right[meta.orderBy!.col]) * direction);
+		return (left, right) => compareValues(left[column], right[column]) * direction;
 	}
-	return rows.sort((left, right) => {
-		for (const column of meta.table.primaryKey) {
+	const columns = meta.table.primaryKey as ReadonlyArray<string>;
+	return (left, right) => {
+		for (const column of columns) {
 			const cmp = compareValues(left[column], right[column]);
 			if (cmp !== 0) return cmp;
 		}
 		return 0;
-	});
+	};
+}
+
+function sortRows(meta: KvQueryMeta, rows: Array<Row>): Array<Row> {
+	return rows.sort(rowComparator(meta));
+}
+
+function decodeRowKey(key: string, columns: number): Array<unknown> | null {
+	const segments = key.split(VALUE_SEPARATOR);
+	if (segments.length !== columns) return null;
+	const values: Array<unknown> = [];
+	for (const segment of segments) {
+		try {
+			values.push(decodeValue(JSON.parse(segment)));
+		} catch {
+			return null;
+		}
+	}
+	return values;
+}
+
+function compareKeyValues(left: ReadonlyArray<unknown>, right: ReadonlyArray<unknown>): number {
+	for (let index = 0; index < left.length; index += 1) {
+		const cmp = compareValues(left[index], right[index]);
+		if (cmp !== 0) return cmp;
+	}
+	return 0;
+}
+
+function compareRowToKeyValues(meta: KvQueryMeta, row: Row, values: ReadonlyArray<unknown>): number {
+	return compareKeyValues(
+		(meta.table.primaryKey as ReadonlyArray<string>).map((column) => row[column]),
+		values,
+	);
 }
 
 function whereClauses(meta: KvQueryMeta): ReadonlyArray<WhereExpr<Row>> {
@@ -627,7 +670,42 @@ function decodePageState(pageState: string | null | undefined): PageState {
 	if (!Number.isInteger(decoded.offset) || decoded.offset < 0) {
 		throw new Error('Invalid Postgres KV page state');
 	}
-	return decoded;
+	if (typeof decoded.after !== 'string') return {offset: decoded.offset};
+	return {offset: decoded.offset, after: decoded.after, keyed: decoded.keyed === true};
+}
+
+function pageableSelect(meta: KvQueryMeta, pageSize: number): boolean {
+	return (
+		meta.action === 'select' &&
+		meta.orderBy === undefined &&
+		meta.limit === undefined &&
+		Number.isInteger(pageSize) &&
+		pageSize > 0
+	);
+}
+
+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<PageEntry>, state: PageState): number {
+	if (state.after === undefined) return state.offset;
+	const index = entries.findIndex((entry) => entry.key === state.after);
+	if (index >= 0) return index + 1;
+	const values = decodeRowKey(state.after, (meta.table.primaryKey as ReadonlyArray<string>).length);
+	if (values === null) return state.offset;
+	let start = 0;
+	while (start < entries.length && compareRowToKeyValues(meta, entries[start]!.row, values) <= 0) start += 1;
+	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<Row> | null {
@@ -850,6 +928,11 @@ export class PostgresKvQueryExecutor {
 		options: {pageSize: number; pageState?: string | null},
 	): Promise<{rows: Array<T>; pageState: string | null}> {
 		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);
+			return {rows: page.rows as Array<T>, pageState: page.pageState};
+		}
 		const rows = await this.executeQuery<T, P>(query);
 		const pageRows = rows.slice(state.offset, state.offset + options.pageSize);
 		const nextOffset = state.offset + pageRows.length;
@@ -859,6 +942,113 @@ export class PostgresKvQueryExecutor {
 		};
 	}
 
+	private async selectPage(
+		meta: KvQueryMeta,
+		params: CassandraParams,
+		state: PageState,
+		pageSize: number,
+	): Promise<{rows: Array<Row>; 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<unknown>} {
+		const values: Array<unknown> = [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<boolean> {
+		if (plan.candidates.kind === 'none') return true;
+		const fragments = planFragmentGroups(plan.candidates)[0]!;
+		const columns = (meta.table.primaryKey as ReadonlyArray<string>).length;
+		let after: string | null = null;
+		let previous: Array<unknown> | 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<Row>; 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<StoredRow>(
+			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<string> | 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,
+		plan: QueryPlan,
+		state: PageState,
+		pageSize: number,
+	): Promise<{rows: Array<Row>; pageState: string | null}> {
+		const entries = this.matchingEntries(meta, await this.candidates(meta, plan, this.client), params);
+		const compare = rowComparator(meta);
+		entries.sort((left, right) => compare(left.row, right.row));
+		const start = pageStart(meta, entries, state);
+		const page = entries.slice(start, start + pageSize);
+		const last = page[page.length - 1];
+		const nextOffset = start + page.length;
+		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,
+		};
+	}
+
 	async executeBatch(
 		queries: Array<{query: string; params: object; meta?: KvQueryMeta}>,
 		atomic = true,
@@ -911,11 +1101,24 @@ export class PostgresKvQueryExecutor {
 		return [...byRowKey.values()];
 	}
 
-	private matchingRows(meta: KvQueryMeta, stored: ReadonlyArray<StoredRow>, params: CassandraParams): Array<Row> {
+	private matchingEntries(
+		meta: KvQueryMeta,
+		stored: ReadonlyArray<StoredRow>,
+		params: CassandraParams,
+	): Array<PageEntry> {
 		const required = queryShape(meta).requiredColumns;
-		return stored
-			.map((entry) => (required ? decodeRowColumns(entry.row_data, required) : decodeRow(entry.row_data)))
-			.filter((row) => matchesWhere(row, meta.where as ReadonlyArray<WhereExpr<Row>> | undefined, params));
+		const entries: Array<PageEntry> = [];
+		for (const entry of stored) {
+			const row = required ? decodeRowColumns(entry.row_data, required) : decodeRow(entry.row_data);
+			if (matchesWhere(row, meta.where as ReadonlyArray<WhereExpr<Row>> | undefined, params)) {
+				entries.push({key: entry.row_key, row});
+			}
+		}
+		return entries;
+	}
+
+	private matchingRows(meta: KvQueryMeta, stored: ReadonlyArray<StoredRow>, params: CassandraParams): Array<Row> {
+		return this.matchingEntries(meta, stored, params).map((entry) => entry.row);
 	}
 
 	private async select(