perf(message): drop two reads and a write from message create (#2123)

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

@@ -0,0 +1,128 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import * as BucketUtils from '@fluxer/snowflake/src/SnowflakeBuckets';
+import {afterEach, beforeEach, describe, expect, it} from 'vitest';
+import {type ChannelID, createChannelID, createMessageID, createUserID, type MessageID} from '../../../BrandedTypes';
+import {fetchOne, setCassandraQueryExecutorForTesting} from '../../../database/CassandraQueryExecution';
+import type {CassandraParams, KvQueryMeta, PreparedQuery} from '../../../database/CassandraTypes';
+import type {ChannelStateRow, MessageRow} from '../../../database/types/MessageTypes';
+import {ChannelState} from '../../../Tables';
+import {InMemoryCassandraQueryExecutor} from '../../../test/InMemoryCassandraQueryExecutor';
+import {MessageDataRepository} from './MessageDataRepository';
+
+const FETCH_CHANNEL_STATE = ChannelState.select({where: ChannelState.where.eq('channel_id'), limit: 1});
+
+class RecordingCassandraQueryExecutor {
+	readonly statements: Array<string> = [];
+	readonly batchAtomicity: Array<boolean> = [];
+	private readonly inner = new InMemoryCassandraQueryExecutor();
+
+	async executeQuery<T = Record<string, unknown>, P extends CassandraParams = CassandraParams>(
+		query: PreparedQuery<P>,
+	): Promise<Array<T>> {
+		this.record(query.kvMeta);
+		return this.inner.executeQuery<T>(query);
+	}
+
+	async executeBatch(
+		queries: Array<{query: string; params: object; meta?: KvQueryMeta}>,
+		atomic = true,
+	): Promise<void> {
+		this.batchAtomicity.push(atomic);
+		for (const entry of queries) {
+			this.record(entry.meta);
+		}
+		await this.inner.executeBatch(queries);
+	}
+
+	countStatements(statement: string): number {
+		return this.statements.filter((entry) => entry === statement).length;
+	}
+
+	private record(meta: KvQueryMeta | undefined): void {
+		if (!meta) return;
+		this.statements.push(`${meta.action}:${meta.table.name}`);
+	}
+}
+
+function makeMessageRow(channelId: ChannelID, messageId: MessageID): MessageRow {
+	return {
+		channel_id: channelId,
+		bucket: BucketUtils.makeBucket(messageId),
+		message_id: messageId,
+		author_id: createUserID(7n),
+		type: 0,
+		webhook_id: null,
+		webhook_name: null,
+		webhook_avatar_hash: null,
+		content: 'hello',
+		edited_timestamp: null,
+		pinned_timestamp: null,
+		flags: 0,
+		mention_everyone: false,
+		mention_users: null,
+		mention_roles: null,
+		mention_channels: null,
+		attachments: null,
+		embeds: null,
+		sticker_items: null,
+		message_reference: null,
+		message_snapshots: null,
+		call: null,
+		nsfw_emojis: null,
+		has_reaction: false,
+		version: 1,
+	};
+}
+
+let executor: RecordingCassandraQueryExecutor;
+
+describe('MessageDataRepository.upsertMessage round trips', () => {
+	beforeEach(() => {
+		executor = new RecordingCassandraQueryExecutor();
+		setCassandraQueryExecutorForTesting(executor);
+	});
+	afterEach(() => {
+		setCassandraQueryExecutorForTesting(new InMemoryCassandraQueryExecutor());
+	});
+	it('writes channel_state once and skips the read-before-write when the caller passes no previous row', async () => {
+		const channelId = createChannelID(1710000000000000000n);
+		const messageId = createMessageID(1720000000000000000n);
+		const repository = new MessageDataRepository();
+
+		await repository.upsertMessage(makeMessageRow(channelId, messageId), null);
+
+		expect(executor.countStatements('select:messages')).toBe(0);
+		expect(executor.countStatements('patch:channel_state')).toBe(1);
+		expect(executor.batchAtomicity).toEqual([false]);
+		expect(await loadChannelState(channelId)).toMatchObject({
+			created_bucket: BucketUtils.makeBucket(channelId),
+			has_messages: true,
+			last_message_id: messageId,
+			last_message_bucket: BucketUtils.makeBucket(messageId),
+		});
+	});
+	it('keeps the newer last_message_id while still recording channel state for an older message', async () => {
+		const channelId = createChannelID(1710000000000000000n);
+		const newerMessageId = createMessageID(1730000000000000000n);
+		const olderMessageId = createMessageID(1720000000000000000n);
+		const repository = new MessageDataRepository();
+		await repository.upsertMessage(makeMessageRow(channelId, newerMessageId), null);
+		executor.statements.length = 0;
+		executor.batchAtomicity.length = 0;
+
+		await repository.upsertMessage(makeMessageRow(channelId, olderMessageId), null);
+
+		expect(executor.countStatements('patch:channel_state')).toBe(1);
+		expect(await loadChannelState(channelId)).toMatchObject({
+			created_bucket: BucketUtils.makeBucket(channelId),
+			has_messages: true,
+			last_message_id: newerMessageId,
+			last_message_bucket: BucketUtils.makeBucket(newerMessageId),
+		});
+	});
+});
+
+async function loadChannelState(channelId: ChannelID): Promise<ChannelStateRow | null> {
+	return fetchOne<ChannelStateRow>(FETCH_CHANNEL_STATE.bind({channel_id: channelId}));
+}

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

@@ -538,6 +538,28 @@ export class MessageDataRepository {
 		);
 	}
 
+	private async recordChannelStateForMessage(
+		channelId: ChannelID,
+		messageId: MessageID,
+		messageBucket: number,
+	): Promise<void> {
+		const state = await this.getChannelState(channelId);
+		const prev = state?.last_message_id ?? null;
+		const patch = {
+			created_bucket: Db.set(BucketUtils.makeBucket(channelId)),
+			has_messages: Db.set(true),
+			updated_at: Db.set(new Date()),
+		};
+		await upsertOne(
+			ChannelState.patchByPk(
+				{channel_id: channelId},
+				prev !== null && messageId <= prev
+					? patch
+					: {...patch, last_message_id: Db.set(messageId), last_message_bucket: Db.set(messageBucket)},
+			),
+		);
+	}
+
 	private async getChannelState(channelId: ChannelID): Promise<ChannelStateRow | null> {
 		return fetchOne<ChannelStateRow>(FETCH_CHANNEL_STATE.bind({channel_id: channelId}));
 	}
@@ -718,19 +740,8 @@ export class MessageDataRepository {
 				updated_at: new Date(),
 			}),
 		);
-		const createdBucket = BucketUtils.makeBucket(data.channel_id);
-		batch.addPrepared(
-			ChannelState.patchByPk(
-				{channel_id: data.channel_id},
-				{
-					created_bucket: Db.set(createdBucket),
-					has_messages: Db.set(true),
-					updated_at: Db.set(new Date()),
-				},
-			),
-		);
-		await batch.execute();
-		await this.advanceChannelStateLastMessageIfNewer(data.channel_id, data.message_id, data.bucket);
+		await batch.execute(false);
+		await this.recordChannelStateForMessage(data.channel_id, data.message_id, data.bucket);
 		return new Message({...data, version: finalVersion});
 	}
 

fluxer_api/src/api/channel/services/message/MessagePersistenceService.ts

@@ -254,7 +254,7 @@ export class MessagePersistenceService {
 			has_reaction: false,
 			version: 1,
 		};
-		const message = await this.channelRepository.messages.upsertMessage(messageRowData);
+		const message = await this.channelRepository.messages.upsertMessage(messageRowData, null);
 		const enqueueDeferredEmbeds = await this.runPostPersistenceOperations({
 			message,
 			params,