perf(read-state): read acked read states in one query (#2196)

fluxer_api/src/api/read_state/ReadStateRepository.test.ts

@@ -2,8 +2,10 @@
 
 import {afterEach, beforeEach, describe, expect, it} from 'vitest';
 import {type ChannelID, createChannelID, createMessageID, createUserID, type UserID} from '../BrandedTypes';
+import {getKvMeta} from '../database/CassandraMetaRegistry';
 import {fetchOne, setCassandraQueryExecutorForTesting, upsertOne} from '../database/CassandraQueryExecution';
 import {defineTable} from '../database/CassandraTableDsl';
+import type {CassandraParams, KvQueryMeta, PreparedQuery} from '../database/CassandraTypes';
 import type {ReadStateRow} from '../database/types/ChannelTypes';
 import {READ_STATE_COLUMNS} from '../database/types/ChannelTypes';
 import {InMemoryCassandraQueryExecutor} from '../test/InMemoryCassandraQueryExecutor';
@@ -19,6 +21,34 @@ const FETCH_READ_STATE = ReadStates.selectCql({
 	limit: 1,
 });
 
+class RecordingCassandraQueryExecutor {
+	readonly statements: Array<string> = [];
+	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 ?? getKvMeta(query.cql));
+		return this.inner.executeQuery<T>(query);
+	}
+
+	async executeBatch(queries: Array<{query: string; params: object; meta?: KvQueryMeta}>): Promise<void> {
+		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}`);
+	}
+}
+
 let executor: InMemoryCassandraQueryExecutor;
 
 describe('ReadStateRepository row storage', () => {
@@ -101,6 +131,59 @@ describe('ReadStateRepository row storage', () => {
 			mention_count: 0,
 		});
 	});
+	it('reads every acked channel of a bulk acknowledgement in one query', async () => {
+		const recording = new RecordingCassandraQueryExecutor();
+		setCassandraQueryExecutorForTesting(recording);
+		const userId = createUserID(1n);
+		const seededChannelIds = [createChannelID(10n), createChannelID(20n), createChannelID(30n)];
+		for (const channelId of seededChannelIds) {
+			await seedReadState(userId, channelId, 100n, 5);
+		}
+		const missingChannelId = createChannelID(40n);
+		recording.statements.length = 0;
+		const repository = new ReadStateRepository();
+		const readStates = await repository.bulkAckMessages(userId, [
+			{channelId: seededChannelIds[0]!, messageId: createMessageID(90n)},
+			{channelId: seededChannelIds[1]!, messageId: createMessageID(200n)},
+			{channelId: seededChannelIds[2]!, messageId: createMessageID(300n)},
+			{channelId: missingChannelId, messageId: createMessageID(400n)},
+		]);
+		expect(recording.countStatements('select:read_states')).toBe(1);
+		expect(readStates.map((state) => state.channelId)).toEqual([...seededChannelIds, missingChannelId]);
+		expect(readStates.map((state) => state.lastMessageId)).toEqual([
+			createMessageID(100n),
+			createMessageID(200n),
+			createMessageID(300n),
+			createMessageID(400n),
+		]);
+		expect(readStates.map((state) => state.mentionCount)).toEqual([5, 0, 0, 0]);
+		expect(await loadReadState(userId, missingChannelId)).toMatchObject({
+			message_id: createMessageID(400n),
+			mention_count: 0,
+		});
+	});
+	it('keeps the last pin timestamp of acked rows and leaves missing rows without one', async () => {
+		const userId = createUserID(1n);
+		const pinnedChannelId = createChannelID(10n);
+		const missingChannelId = createChannelID(20n);
+		const lastPinTimestamp = new Date('2026-01-01T00:00:00.000Z');
+		await upsertOne(
+			ReadStates.upsertAll({
+				user_id: userId,
+				channel_id: pinnedChannelId,
+				message_id: createMessageID(100n),
+				mention_count: 2,
+				last_pin_timestamp: lastPinTimestamp,
+			}),
+		);
+		const repository = new ReadStateRepository();
+		const readStates = await repository.bulkAckMessages(userId, [
+			{channelId: pinnedChannelId, messageId: createMessageID(200n)},
+			{channelId: missingChannelId, messageId: createMessageID(200n)},
+		]);
+		expect(readStates[0]?.lastPinTimestamp).toEqual(lastPinTimestamp);
+		expect(readStates[1]?.lastPinTimestamp).toBeNull();
+	});
 });
 
 async function seedReadState(

fluxer_api/src/api/read_state/ReadStateRepository.ts

@@ -2,7 +2,14 @@
 
 import type {ChannelID, MessageID, UserID} from '../BrandedTypes';
 import {channelIdToMessageId} from '../BrandedTypes';
-import {BatchBuilder, deleteOneOrMany, fetchMany, fetchOne, upsertOne} from '../database/CassandraQueryExecution';
+import {
+	BatchBuilder,
+	deleteOneOrMany,
+	fetchMany,
+	fetchManyInChunks,
+	fetchOne,
+	upsertOne,
+} from '../database/CassandraQueryExecution';
 import {defineTable} from '../database/CassandraTableDsl';
 import {Db, type DbOp} from '../database/CassandraTypes';
 import type {ReadStateRow} from '../database/types/ChannelTypes';
@@ -22,6 +29,9 @@ const FETCH_READ_STATE_BY_USER_AND_CHANNEL_CQL = ReadStates.selectCql({
 	where: [ReadStates.where.eq('user_id'), ReadStates.where.eq('channel_id')],
 	limit: 1,
 });
+const FETCH_READ_STATES_BY_CHANNEL_IDS_CQL = ReadStates.selectCql({
+	where: [ReadStates.where.eq('user_id'), ReadStates.where.in('channel_id', 'channel_ids')],
+});
 const BULK_READ_STATE_BATCH_QUERY_LIMIT = 50;
 
 export class ReadStateRepository implements IReadStateRepository {
@@ -215,18 +225,16 @@ export class ReadStateRepository implements IReadStateRepository {
 			messageId: MessageID;
 		}>,
 	): Promise<Array<ReadState>> {
-		const currentRows = await Promise.all(
-			readStates.map((readState) =>
-				fetchOne<ReadStateRow>(FETCH_READ_STATE_BY_USER_AND_CHANNEL_CQL, {
-					user_id: userId,
-					channel_id: readState.channelId,
-				}),
-			),
+		const currentRows = await fetchManyInChunks<ReadStateRow, ChannelID>(
+			FETCH_READ_STATES_BY_CHANNEL_IDS_CQL,
+			readStates.map((readState) => readState.channelId),
+			(chunk) => ({user_id: userId, channel_ids: chunk}),
 		);
+		const currentRowsByChannel = new Map(currentRows.map((row) => [row.channel_id, row]));
 		const batch = new BatchBuilder();
 		const results: Array<ReadState> = [];
-		for (const [index, readState] of readStates.entries()) {
-			const currentReadState = currentRows[index];
+		for (const readState of readStates) {
+			const currentReadState = currentRowsByChannel.get(readState.channelId);
 			if (currentReadState?.message_id != null && currentReadState.message_id > readState.messageId) {
 				results.push(new ReadState(currentReadState));
 				continue;