perf(cassandra): skip re-registering identical select metadata (#2200)

fluxer_api/src/api/channel/repositories/message/MessageDataRepository.ts

@@ -4,7 +4,7 @@ import {generateSnowflake} from '@fluxer/snowflake/src/Snowflake';
 import * as BucketUtils from '@fluxer/snowflake/src/SnowflakeBuckets';
 import type {ChannelID, MessageID} from '../../../BrandedTypes';
 import {BatchBuilder, deleteOneOrMany, fetchMany, fetchOne, upsertOne} from '../../../database/CassandraQueryExecution';
-import {Db} from '../../../database/CassandraTypes';
+import {Db, type QueryTemplate} from '../../../database/CassandraTypes';
 import {buildPatchFromData, executeVersionedUpdate} from '../../../database/CassandraVersionedUpdate';
 import type {ChannelMessageBucketRow, ChannelStateRow, MessageRow} from '../../../database/types/MessageTypes';
 import {MESSAGE_COLUMNS} from '../../../database/types/MessageTypes';
@@ -38,6 +38,77 @@ const FETCH_CHANNEL_STATE = ChannelState.select({
 	limit: 1,
 });
 
+function memoizeQueryByLimit(build: (limit: number) => QueryTemplate): (limit: number) => QueryTemplate {
+	const cache = new Map<number, QueryTemplate>();
+	return (limit: number) => {
+		const cached = cache.get(limit);
+		if (cached) return cached;
+		const query = build(limit);
+		cache.set(limit, query);
+		return query;
+	};
+}
+
+const fetchMessagesBeforeQuery = memoizeQueryByLimit((limit) =>
+	Messages.select({
+		where: [
+			Messages.where.eq('channel_id'),
+			Messages.where.eq('bucket'),
+			Messages.where.lt('message_id', 'before_message_id'),
+		],
+		orderBy: {col: 'message_id', direction: 'DESC'},
+		limit,
+	}),
+);
+const fetchMessagesAfterDescQuery = memoizeQueryByLimit((limit) =>
+	Messages.select({
+		where: [
+			Messages.where.eq('channel_id'),
+			Messages.where.eq('bucket'),
+			Messages.where.gt('message_id', 'after_message_id'),
+		],
+		orderBy: {col: 'message_id', direction: 'DESC'},
+		limit,
+	}),
+);
+const fetchMessagesBetweenQuery = memoizeQueryByLimit((limit) =>
+	Messages.select({
+		where: [
+			Messages.where.eq('channel_id'),
+			Messages.where.eq('bucket'),
+			Messages.where.gt('message_id', 'after_message_id'),
+			Messages.where.lt('message_id', 'before_message_id'),
+		],
+		orderBy: {col: 'message_id', direction: 'DESC'},
+		limit,
+	}),
+);
+const fetchMessagesLatestDescQuery = memoizeQueryByLimit((limit) =>
+	Messages.select({
+		where: [Messages.where.eq('channel_id'), Messages.where.eq('bucket')],
+		orderBy: {col: 'message_id', direction: 'DESC'},
+		limit,
+	}),
+);
+const fetchMessagesAfterAscQuery = memoizeQueryByLimit((limit) =>
+	Messages.select({
+		where: [
+			Messages.where.eq('channel_id'),
+			Messages.where.eq('bucket'),
+			Messages.where.gt('message_id', 'after_message_id'),
+		],
+		orderBy: {col: 'message_id', direction: 'ASC'},
+		limit,
+	}),
+);
+const fetchMessagesOldestAscQuery = memoizeQueryByLimit((limit) =>
+	Messages.select({
+		where: [Messages.where.eq('channel_id'), Messages.where.eq('bucket')],
+		orderBy: {col: 'message_id', direction: 'ASC'},
+		limit,
+	}),
+);
+
 export class MessageDataRepository {
 	async listMessages(
 		channelId: ChannelID,
@@ -68,71 +139,6 @@ export class MessageDataRepository {
 		return this.listMessagesLatest(channelId, limit);
 	}
 
-	private makeFetchMessagesBefore(limit: number) {
-		return Messages.select({
-			where: [
-				Messages.where.eq('channel_id'),
-				Messages.where.eq('bucket'),
-				Messages.where.lt('message_id', 'before_message_id'),
-			],
-			orderBy: {col: 'message_id', direction: 'DESC'},
-			limit,
-		});
-	}
-
-	private makeFetchMessagesAfterDesc(limit: number) {
-		return Messages.select({
-			where: [
-				Messages.where.eq('channel_id'),
-				Messages.where.eq('bucket'),
-				Messages.where.gt('message_id', 'after_message_id'),
-			],
-			orderBy: {col: 'message_id', direction: 'DESC'},
-			limit,
-		});
-	}
-
-	private makeFetchMessagesBetween(limit: number) {
-		return Messages.select({
-			where: [
-				Messages.where.eq('channel_id'),
-				Messages.where.eq('bucket'),
-				Messages.where.gt('message_id', 'after_message_id'),
-				Messages.where.lt('message_id', 'before_message_id'),
-			],
-			orderBy: {col: 'message_id', direction: 'DESC'},
-			limit,
-		});
-	}
-
-	private makeFetchMessagesLatestDesc(limit: number) {
-		return Messages.select({
-			where: [Messages.where.eq('channel_id'), Messages.where.eq('bucket')],
-			orderBy: {col: 'message_id', direction: 'DESC'},
-			limit,
-		});
-	}
-
-	private makeFetchMessagesAfterAsc(limit: number) {
-		return Messages.select({
-			where: [
-				Messages.where.eq('channel_id'),
-				Messages.where.eq('bucket'),
-				Messages.where.gt('message_id', 'after_message_id'),
-			],
-			orderBy: {col: 'message_id', direction: 'ASC'},
-			limit,
-		});
-	}
-
-	private makeFetchMessagesOldestAsc(limit: number) {
-		return Messages.select({
-			where: [Messages.where.eq('channel_id'), Messages.where.eq('bucket')],
-			orderBy: {col: 'message_id', direction: 'ASC'},
-			limit,
-		});
-	}
-
 	private async listMessagesLatest(channelId: ChannelID, limit: number): Promise<Array<Message>> {
 		const state = await this.getChannelState(channelId);
 		const nowId = generateSnowflake();
@@ -387,7 +393,7 @@ export class MessageDataRepository {
 			'fetchRowsForBucketAsc parameters',
 		);
 		if (bucket === meta.afterBucket) {
-			const q = this.makeFetchMessagesAfterAsc(limit);
+			const q = fetchMessagesAfterAscQuery(limit);
 			const rows = await fetchMany<MessageRow>(
 				q.bind({
 					channel_id: channelId,
@@ -397,7 +403,7 @@ export class MessageDataRepository {
 			);
 			return {rows, unbounded: false};
 		}
-		const q = this.makeFetchMessagesOldestAsc(limit);
+		const q = fetchMessagesOldestAscQuery(limit);
 		const rows = await fetchMany<MessageRow>(q.bind({channel_id: channelId, bucket}));
 		return {rows, unbounded: true};
 	}
@@ -431,7 +437,7 @@ export class MessageDataRepository {
 			'fetchRowsForBucket parameters',
 		);
 		if (meta.before && meta.after && meta.beforeBucket === bucket && meta.afterBucket === bucket) {
-			const q = this.makeFetchMessagesBetween(limit);
+			const q = fetchMessagesBetweenQuery(limit);
 			const rows = await fetchMany<MessageRow>(
 				q.bind({
 					channel_id: channelId,
@@ -443,7 +449,7 @@ export class MessageDataRepository {
 			return {rows, unbounded: false};
 		}
 		if (meta.before && meta.beforeBucket === bucket) {
-			const q = this.makeFetchMessagesBefore(limit);
+			const q = fetchMessagesBeforeQuery(limit);
 			const rows = await fetchMany<MessageRow>(
 				q.bind({
 					channel_id: channelId,
@@ -454,7 +460,7 @@ export class MessageDataRepository {
 			return {rows, unbounded: false};
 		}
 		if (meta.after && meta.afterBucket === bucket) {
-			const q = this.makeFetchMessagesAfterDesc(limit);
+			const q = fetchMessagesAfterDescQuery(limit);
 			const rows = await fetchMany<MessageRow>(
 				q.bind({
 					channel_id: channelId,
@@ -464,7 +470,7 @@ export class MessageDataRepository {
 			);
 			return {rows, unbounded: false};
 		}
-		const q = this.makeFetchMessagesLatestDesc(limit);
+		const q = fetchMessagesLatestDescQuery(limit);
 		const rows = await fetchMany<MessageRow>(q.bind({channel_id: channelId, bucket}));
 		return {rows, unbounded: true};
 	}

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

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import {describe, expect, it} from 'vitest';
+import {getKvMeta, registerKvMeta} from './CassandraMetaRegistry';
+import type {KvQueryMeta, KvTableSpec} from './CassandraTypes';
+
+const MetaRegistryTestRows: KvTableSpec = {
+	name: 'meta_registry_test_rows',
+	columns: ['id', 'value'],
+	primaryKey: ['id'],
+	partitionKey: ['id'],
+};
+
+function selectMeta(limit: number): KvQueryMeta {
+	return {action: 'select', table: MetaRegistryTestRows, limit};
+}
+
+describe('CassandraMetaRegistry', () => {
+	it('resolves lookups whose whitespace differs from the registered statement', () => {
+		const meta = selectMeta(10);
+		registerKvMeta('SELECT id, value FROM meta_registry_test_rows\n\tWHERE id = :id;', meta);
+		expect(getKvMeta('SELECT id, value FROM meta_registry_test_rows WHERE id = :id;')).toBe(meta);
+		expect(getKvMeta('  SELECT id,   value FROM meta_registry_test_rows\n WHERE id = :id;  ')).toBe(meta);
+	});
+
+	it('keeps the most recently registered metadata for a repeated statement', () => {
+		const cql = 'SELECT id, value FROM meta_registry_test_rows WHERE id = :id LIMIT :limit_bind;';
+		const first = selectMeta(10);
+		const second = selectMeta(20);
+		registerKvMeta(cql, first);
+		expect(getKvMeta(cql)).toBe(first);
+		registerKvMeta(cql, second);
+		expect(getKvMeta(cql)).toBe(second);
+		expect(getKvMeta(`  ${cql.replace(/ /g, '   ')}\n`)).toBe(second);
+	});
+});

fluxer_api/src/api/database/CassandraMetaRegistry.ts

@@ -15,6 +15,7 @@ interface TableMetadata {
 }
 
 const kvMetaRegistry = new Map<string, KvQueryMeta<Record<string, unknown>>>();
+const kvMetaKeyCache = new Map<string, string>();
 const tableRegistry = new Map<string, TableMetadata>();
 
 export function registerTableSpec<Row extends object>(tableSpec: KvTableSpec<Row>): void {
@@ -32,10 +33,17 @@ function normalizeCqlForRegistry(cql: string): string {
 }
 
 export function registerKvMeta(cql: string, meta: KvQueryMeta): void {
-	kvMetaRegistry.set(normalizeCqlForRegistry(cql), meta);
+	let key = kvMetaKeyCache.get(cql);
+	if (key === undefined) {
+		key = normalizeCqlForRegistry(cql);
+		kvMetaKeyCache.set(cql, key);
+	}
+	kvMetaRegistry.set(key, meta);
 }
 
 export function getKvMeta(cql: string): KvQueryMeta | undefined {
+	const key = kvMetaKeyCache.get(cql);
+	if (key !== undefined) return kvMetaRegistry.get(key);
 	return kvMetaRegistry.get(normalizeCqlForRegistry(cql));
 }