fix(kv): align rust and typescript schema migration (#2328)

fluxer_api/src/api/database/PostgresKvQueryExecutor.ts

@@ -73,6 +73,9 @@ 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 POSTGRES_KV_SCHEMA_ATTEMPTS = 3;
+const POSTGRES_KV_SCHEMA_RETRY_DELAY_MS = 250;
+const POSTGRES_KV_CONCURRENT_DDL_CODES = new Set(['23505', '42P07', '42710']);
 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()';
@@ -816,7 +819,7 @@ WHERE att.attrelid = to_regclass($1)
 	return result.rows[0]?.c_collated === true;
 }
 
-export async function ensurePostgresKvSchema(client: IPostgresClient): Promise<void> {
+async function ensurePostgresKvSchemaOnce(client: IPostgresClient): Promise<void> {
 	const kvTable = client.kvTable();
 	const table = quoteIdentifier(kvTable);
 	await client.transaction(async (db) => {
@@ -885,6 +888,28 @@ ON CONFLICT (table_name, row_key) DO NOTHING`,
 	});
 }
 
+function isConcurrentDdlConflict(error: unknown): boolean {
+	return (
+		typeof error === 'object' &&
+		error !== null &&
+		'code' in error &&
+		POSTGRES_KV_CONCURRENT_DDL_CODES.has(String((error as {code: unknown}).code))
+	);
+}
+
+export async function ensurePostgresKvSchema(client: IPostgresClient): Promise<void> {
+	for (let attempt = 1; attempt <= POSTGRES_KV_SCHEMA_ATTEMPTS; attempt += 1) {
+		try {
+			await ensurePostgresKvSchemaOnce(client);
+			return;
+		} catch (error) {
+			if (attempt === POSTGRES_KV_SCHEMA_ATTEMPTS || !isConcurrentDdlConflict(error)) throw error;
+			logWarn({table: client.kvTable(), attempt}, 'Postgres KV schema hit a concurrent DDL conflict, retrying');
+			await new Promise((resolve) => setTimeout(resolve, POSTGRES_KV_SCHEMA_RETRY_DELAY_MS));
+		}
+	}
+}
+
 export async function pruneExpiredPostgresKvRows(client: IPostgresClient, batchSize = 5000): Promise<number> {
 	if (!Number.isInteger(batchSize) || batchSize <= 0) {
 		throw new Error('Postgres KV prune batch size must be a positive integer');

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

@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import fs from 'node:fs';
+import path from 'node:path';
+import {fileURLToPath} from 'node:url';
+import {describe, expect, it} from 'vitest';
+
+const THIS_DIR = path.dirname(fileURLToPath(import.meta.url));
+const REPO_ROOT = path.resolve(THIS_DIR, '../../../..');
+
+const TYPESCRIPT_SOURCE = fs.readFileSync(path.join(THIS_DIR, 'PostgresKvQueryExecutor.ts'), 'utf8');
+const RUST_SOURCE = fs.readFileSync(path.join(REPO_ROOT, 'fluxer_svc/src/postgres.rs'), 'utf8');
+
+const RUST_INDEX_SUFFIXES = [
+	'expires_idx',
+	'message_reactions_message_idx',
+	'messages_message_idx',
+	'partition_idx',
+	'partition_row_idx',
+	'row_key_c_idx',
+];
+
+const SHARED_LITERALS = ['120s', '30min', '__fluxer_schema_migrations', 'messages_partition_key_v1'];
+
+const BACKFILL_FRAGMENTS = [
+	"split_part(row_key, chr(31), 3) <> ''",
+	'split_part(row_key, chr(31), 1) || chr(31) || split_part(row_key, chr(31), 2)',
+];
+
+function indexSuffixes(source: string, pattern: RegExp): Array<string> {
+	return [...new Set([...source.matchAll(pattern)].map((match) => match[1]!))].sort();
+}
+
+function lockNamespace(source: string, pattern: RegExp): number {
+	const match = pattern.exec(source);
+	expect(match, 'missing Postgres KV schema lock namespace').not.toBeNull();
+	return Number.parseInt(match![1]!.replaceAll('_', ''), 16);
+}
+
+describe('Postgres KV schema parity between the API and fluxer_svc', () => {
+	it('creates the same indexes apart from the API-only numeric index', () => {
+		const typescript = indexSuffixes(TYPESCRIPT_SOURCE, /\$\{kvTable\}_([a-z_]+)/g);
+		const rust = indexSuffixes(RUST_SOURCE, /\{kv_table\}_([a-z_]+)/g);
+		expect(rust).toEqual(RUST_INDEX_SUFFIXES);
+		expect(typescript.filter((suffix) => !rust.includes(suffix))).toEqual(['row_key_numeric_idx']);
+		expect(rust.filter((suffix) => !typescript.includes(suffix))).toEqual([]);
+	});
+
+	it('takes the same advisory lock namespace', () => {
+		expect(lockNamespace(RUST_SOURCE, /POSTGRES_KV_SCHEMA_LOCK_NAMESPACE: i32 = (0x[0-9a-fA-F_]+)/)).toBe(
+			lockNamespace(TYPESCRIPT_SOURCE, /POSTGRES_KV_SCHEMA_LOCK_NAMESPACE = (0x[0-9a-fA-F_]+)/),
+		);
+	});
+
+	it('shares the schema timeouts and the migration marker identity', () => {
+		for (const literal of SHARED_LITERALS) {
+			expect(TYPESCRIPT_SOURCE, literal).toContain(literal);
+			expect(RUST_SOURCE, literal).toContain(literal);
+		}
+	});
+
+	it('never disables the statement timeout in fluxer_svc', () => {
+		expect(RUST_SOURCE).not.toContain("set_config('statement_timeout', '0'");
+	});
+
+	it('backfills message partition keys with byte-identical SQL', () => {
+		for (const fragment of BACKFILL_FRAGMENTS) {
+			expect(TYPESCRIPT_SOURCE, fragment).toContain(fragment);
+			expect(RUST_SOURCE, fragment).toContain(fragment);
+		}
+	});
+});

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

@@ -367,6 +367,37 @@ suite('postgres kv upgrade safety', () => {
 		]);
 	}, 120_000);
 
+	it('survives a peer that creates the table without the schema lock', async () => {
+		const RACE = `${KV}_race`;
+		await raw.query(`DROP TABLE IF EXISTS ${RACE}`);
+		let release = () => {};
+		const gate = new Promise<void>((resolve) => {
+			release = resolve;
+		});
+		const peer = raw.transaction(async (db) => {
+			await db.query(`
+CREATE TABLE IF NOT EXISTS ${RACE} (
+	table_name text NOT NULL,
+	partition_key text COLLATE "C" NOT NULL,
+	row_key text COLLATE "C" NOT NULL,
+	row_data jsonb NOT NULL,
+	expires_at timestamptz,
+	updated_at timestamptz NOT NULL DEFAULT now(),
+	PRIMARY KEY (table_name, row_key)
+)`);
+			await gate;
+		});
+		const booting = ensurePostgresKvSchema(new TableClient(raw, RACE)).then(
+			() => 'ok',
+			(error: Error) => `failed: ${error.message}`,
+		);
+		await sleep(500);
+		release();
+		await peer;
+		expect(await booting).toBe('ok');
+		await raw.query(`DROP TABLE IF EXISTS ${RACE}`);
+	}, 120_000);
+
 	it('backfills the messages partition key once and never scans for it again', async () => {
 		const BACKFILL = 'kv_backfill';
 		const SEP = String.fromCharCode(31);