perf(api): cache channel and guild reads for the request (#2220)

fluxer_api/src/api/auth/AuthLogin.ts

@@ -29,7 +29,7 @@ import type {KVAccountDeletionQueueService} from '../infrastructure/KVAccountDel
 import {REGISTRATION_PENDING_APPROVAL_TRAIT, REGISTRATION_REJECTED_TRAIT} from '../instance/InstanceConfigRepository';
 import type {InviteService} from '../invite/InviteService';
 import {Logger} from '../Logger';
-import type {RequestCache} from '../middleware/RequestCacheMiddleware';
+import {createRequestCache} from '../middleware/RequestCacheMiddleware';
 import {getInstanceConfigRepository} from '../middleware/ServiceSingletons';
 import type {User} from '../models/User';
 import {lookupGeoip} from '../utils/IpUtils';
@@ -38,19 +38,6 @@ import * as AuthPassword from './AuthPassword';
 import * as AuthSession from './AuthSession';
 import * as AuthUtility from './AuthUtility';
 
-function createRequestCache(): RequestCache {
-	const userPartials = new Map();
-	const messageMentionChannels = new Map();
-	return {
-		userPartials,
-		messageMentionChannels,
-		clear: () => {
-			userPartials.clear();
-			messageMentionChannels.clear();
-		},
-	};
-}
-
 interface LoginParams {
 	data: LoginRequest;
 	request: Request;

fluxer_api/src/api/channel/ChannelRepository.ts

@@ -3,6 +3,7 @@
 import type {AttachmentID, ChannelID, EmojiID, GuildID, MessageID, UserID} from '../BrandedTypes';
 import type {ChannelRow} from '../database/types/ChannelTypes';
 import type {MessageRow} from '../database/types/MessageTypes';
+import type {RequestCache} from '../middleware/RequestCacheMiddleware';
 import type {Channel} from '../models/Channel';
 import type {Message} from '../models/Message';
 import type {MessageReaction} from '../models/MessageReaction';
@@ -12,9 +13,9 @@ import {ChannelRepository as NewChannelRepository} from './repositories/ChannelR
 export class ChannelRepository extends IChannelRepository {
 	private repository: NewChannelRepository;
 
-	constructor() {
+	constructor(requestCache?: RequestCache) {
 		super();
-		this.repository = new NewChannelRepository();
+		this.repository = new NewChannelRepository(requestCache);
 	}
 
 	get channelData() {

fluxer_api/src/api/channel/repositories/ChannelDataRepository.test.ts

@@ -0,0 +1,135 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import {ChannelTypes} from '@fluxer/constants/src/ChannelConstants';
+import {afterEach, beforeEach, describe, expect, it} from 'vitest';
+import {type ChannelID, createChannelID, createGuildID, createMessageID} from '../../BrandedTypes';
+import {setCassandraQueryExecutorForTesting, upsertOne} from '../../database/CassandraQueryExecution';
+import type {CassandraParams, KvQueryMeta, PreparedQuery} from '../../database/CassandraTypes';
+import type {ChannelRow} from '../../database/types/ChannelTypes';
+import {createRequestCache} from '../../middleware/RequestCacheMiddleware';
+import {Channel} from '../../models/Channel';
+import {Channels} from '../../Tables';
+import {InMemoryCassandraQueryExecutor} from '../../test/InMemoryCassandraQueryExecutor';
+import {ChannelDataRepository} from './ChannelDataRepository';
+
+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);
+		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}`);
+	}
+}
+
+function makeChannelRow(channelId: ChannelID): ChannelRow {
+	return {
+		channel_id: channelId,
+		guild_id: createGuildID(1810000000000000000n),
+		type: ChannelTypes.GUILD_TEXT,
+		name: 'general',
+		topic: null,
+		icon_hash: null,
+		url: null,
+		parent_id: null,
+		position: 0,
+		owner_id: null,
+		recipient_ids: null,
+		nsfw: false,
+		content_warning_level: null,
+		content_warning_text: null,
+		rate_limit_per_user: 0,
+		bitrate: null,
+		user_limit: null,
+		voice_connection_limit: null,
+		rtc_region: null,
+		last_message_id: null,
+		last_pin_timestamp: null,
+		permission_overwrites: null,
+		nicks: null,
+		soft_deleted: false,
+		indexed_at: null,
+		version: 1,
+	};
+}
+
+let executor: RecordingCassandraQueryExecutor;
+
+describe('ChannelDataRepository request cache', () => {
+	beforeEach(() => {
+		executor = new RecordingCassandraQueryExecutor();
+		setCassandraQueryExecutorForTesting(executor);
+	});
+	afterEach(() => {
+		setCassandraQueryExecutorForTesting(new InMemoryCassandraQueryExecutor());
+	});
+	it('reads the channel row once per lookup without a request cache', async () => {
+		const channelId = createChannelID(1910000000000000000n);
+		await upsertOne(Channels.upsertAll(makeChannelRow(channelId)));
+		executor.statements.length = 0;
+		const repository = new ChannelDataRepository();
+
+		await repository.findUnique(channelId);
+		await repository.findUnique(channelId);
+
+		expect(executor.countStatements('select:channels')).toBe(2);
+	});
+	it('serves the prefetched channel once and falls back to the database afterwards', async () => {
+		const channelId = createChannelID(1910000000000000000n);
+		const row = makeChannelRow(channelId);
+		await upsertOne(Channels.upsertAll(row));
+		executor.statements.length = 0;
+		const requestCache = createRequestCache();
+		requestCache.channels.set(channelId, new Channel(row));
+		const repository = new ChannelDataRepository(requestCache);
+
+		const prefetched = await repository.findUnique(channelId);
+		expect(prefetched?.id).toBe(channelId);
+		expect(executor.countStatements('select:channels')).toBe(0);
+
+		const reread = await repository.findUnique(channelId);
+		expect(reread?.id).toBe(channelId);
+		expect(executor.countStatements('select:channels')).toBe(1);
+	});
+	it('keeps a prefetched miss distinguishable from an absent entry', async () => {
+		const channelId = createChannelID(1920000000000000000n);
+		const requestCache = createRequestCache();
+		requestCache.channels.set(channelId, null);
+		const repository = new ChannelDataRepository(requestCache);
+
+		expect(await repository.findUnique(channelId)).toBeNull();
+		expect(executor.countStatements('select:channels')).toBe(0);
+	});
+	it('drops the prefetched channel when the same request writes it', async () => {
+		const channelId = createChannelID(1930000000000000000n);
+		const row = makeChannelRow(channelId);
+		await upsertOne(Channels.upsertAll(row));
+		executor.statements.length = 0;
+		const requestCache = createRequestCache();
+		requestCache.channels.set(channelId, null);
+		const repository = new ChannelDataRepository(requestCache);
+
+		await repository.updateLastMessageId(channelId, createMessageID(1940000000000000000n));
+
+		expect(requestCache.channels.has(channelId)).toBe(false);
+		expect(await repository.findUnique(channelId)).not.toBeNull();
+	});
+});

fluxer_api/src/api/channel/repositories/ChannelDataRepository.ts

@@ -7,6 +7,7 @@ import {buildPatchFromData, executeVersionedUpdate} from '../../database/Cassand
 import type {ChannelRow} from '../../database/types/ChannelTypes';
 import {CHANNEL_COLUMNS} from '../../database/types/ChannelTypes';
 import {Logger} from '../../Logger';
+import type {RequestCache} from '../../middleware/RequestCacheMiddleware';
 import {Channel} from '../../models/Channel';
 import {Channels, ChannelsByGuild, PrivateChannels} from '../../Tables';
 import {
@@ -33,7 +34,15 @@ const FETCH_OPEN_PRIVATE_CHANNEL_TARGET = PrivateChannels.selectCql({
 });
 
 export class ChannelDataRepository extends IChannelDataRepository {
+	constructor(private readonly requestCache?: RequestCache) {
+		super();
+	}
+
 	async findUnique(channelId: ChannelID): Promise<Channel | null> {
+		const prefetched = this.requestCache?.takeChannel(channelId);
+		if (prefetched !== undefined) {
+			return prefetched;
+		}
 		const channel = await fetchOne<ChannelRow>(
 			FETCH_CHANNEL_BY_ID.bind({
 				channel_id: channelId,
@@ -45,6 +54,7 @@ export class ChannelDataRepository extends IChannelDataRepository {
 
 	async upsert(data: ChannelRow, oldData?: ChannelRow | null): Promise<Channel> {
 		const channelId = data.channel_id;
+		this.requestCache?.channels.delete(channelId);
 		const result = await executeVersionedUpdate<ChannelRow, 'channel_id' | 'soft_deleted'>(
 			async () => fetchOne<ChannelRow>(FETCH_CHANNEL_BY_ID.bind({channel_id: channelId, soft_deleted: false})),
 			(current) => ({
@@ -68,6 +78,7 @@ export class ChannelDataRepository extends IChannelDataRepository {
 	}
 
 	async updateLastMessageId(channelId: ChannelID, messageId: MessageID): Promise<void> {
+		this.requestCache?.channels.delete(channelId);
 		const existing = await fetchOne<ChannelRow>(
 			FETCH_CHANNEL_BY_ID.bind({
 				channel_id: channelId,
@@ -159,6 +170,7 @@ export class ChannelDataRepository extends IChannelDataRepository {
 	}
 
 	async delete(channelId: ChannelID, guildId?: GuildID): Promise<void> {
+		this.requestCache?.channels.delete(channelId);
 		const batch = new BatchBuilder();
 		batch.addPrepared(
 			Channels.deleteByPk({

fluxer_api/src/api/channel/repositories/ChannelRepository.ts

@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
+import type {RequestCache} from '../../middleware/RequestCacheMiddleware';
 import {ChannelDataRepository} from './ChannelDataRepository';
 import {IChannelRepositoryAggregate} from './IChannelRepositoryAggregate';
 import {MessageInteractionRepository} from './MessageInteractionRepository';
@@ -10,9 +11,9 @@ export class ChannelRepository extends IChannelRepositoryAggregate {
 	readonly messages: MessageRepository;
 	readonly messageInteractions: MessageInteractionRepository;
 
-	constructor() {
+	constructor(requestCache?: RequestCache) {
 		super();
-		this.channelData = new ChannelDataRepository();
+		this.channelData = new ChannelDataRepository(requestCache);
 		this.messages = new MessageRepository(this.channelData);
 		this.messageInteractions = new MessageInteractionRepository(this.messages);
 	}

fluxer_api/src/api/guild/repositories/GuildDataRepository.ts

@@ -4,6 +4,7 @@ import type {GuildID, UserID} from '../../BrandedTypes';
 import {BatchBuilder, fetchMany, fetchOne} from '../../database/CassandraQueryExecution';
 import {buildPatchFromData, executeVersionedUpdate} from '../../database/CassandraVersionedUpdate';
 import {GUILD_COLUMNS, type GuildMemberByUserIdRow, type GuildRow} from '../../database/types/GuildTypes';
+import type {RequestCache} from '../../middleware/RequestCacheMiddleware';
 import {Guild} from '../../models/Guild';
 import {GuildMembersByUserId, Guilds} from '../../Tables';
 import {IGuildDataRepository} from './IGuildDataRepository';
@@ -26,7 +27,15 @@ function createFetchAllGuildsFirstPageQuery(limit: number) {
 }
 
 export class GuildDataRepository extends IGuildDataRepository {
+	constructor(private readonly requestCache?: RequestCache) {
+		super();
+	}
+
 	async findUnique(guildId: GuildID): Promise<Guild | null> {
+		const prefetched = this.requestCache?.takeGuild(guildId);
+		if (prefetched !== undefined) {
+			return prefetched;
+		}
 		const guild = await fetchOne<GuildRow>(FETCH_GUILD_BY_ID_QUERY, {
 			guild_id: guildId,
 		});
@@ -87,6 +96,7 @@ export class GuildDataRepository extends IGuildDataRepository {
 
 	async upsert(data: GuildRow, oldData?: GuildRow | null, _previousOwnerId?: UserID): Promise<Guild> {
 		const guildId = data.guild_id;
+		this.requestCache?.guilds.delete(guildId);
 		const result = await executeVersionedUpdate<GuildRow, 'guild_id'>(
 			async () => fetchOne<GuildRow>(FETCH_GUILD_BY_ID_QUERY, {guild_id: guildId}),
 			(current) => ({
@@ -105,6 +115,7 @@ export class GuildDataRepository extends IGuildDataRepository {
 		oldData?: GuildRow | null,
 		_previousOwnerId?: UserID,
 	): Promise<Guild> {
+		this.requestCache?.guilds.delete(guildId);
 		const result = await executeVersionedUpdate<GuildRow, 'guild_id'>(
 			async () => fetchOne<GuildRow>(FETCH_GUILD_BY_ID_QUERY, {guild_id: guildId}),
 			(current) => ({
@@ -122,6 +133,7 @@ export class GuildDataRepository extends IGuildDataRepository {
 	}
 
 	async delete(guildId: GuildID, _ownerId?: UserID): Promise<void> {
+		this.requestCache?.guilds.delete(guildId);
 		const guild = await fetchOne<GuildRow>(FETCH_GUILD_BY_ID_QUERY, {guild_id: guildId});
 		if (!guild) {
 			return;

fluxer_api/src/api/guild/repositories/GuildModerationRepository.ts

@@ -13,6 +13,7 @@ import type {
 	GuildBanRow,
 	GuildRow,
 } from '../../database/types/GuildTypes';
+import type {RequestCache} from '../../middleware/RequestCacheMiddleware';
 import {GuildAuditLog} from '../../models/GuildAuditLog';
 import {GuildBan} from '../../models/GuildBan';
 import {
@@ -55,6 +56,10 @@ const FETCH_GUILD_AUDIT_LOGS_BY_IDS_QUERY = GuildAuditLogs.selectCql({
 });
 
 export class GuildModerationRepository extends IGuildModerationRepository {
+	constructor(private readonly requestCache?: RequestCache) {
+		super();
+	}
+
 	async getBan(guildId: GuildID, userId: UserID): Promise<GuildBan | null> {
 		const ban = await fetchOne<GuildBanRow>(FETCH_GUILD_BAN_BY_GUILD_AND_USER_ID_QUERY, {
 			guild_id: guildId,
@@ -279,6 +284,7 @@ export class GuildModerationRepository extends IGuildModerationRepository {
 	}
 
 	async updateAuditLogsIndexedAt(guildId: GuildID, indexedAt: Date | null): Promise<void> {
+		this.requestCache?.guilds.delete(guildId);
 		await executeVersionedUpdate<GuildRow, 'guild_id'>(
 			() => fetchOne<GuildRow>(FETCH_GUILD_BY_ID_QUERY, {guild_id: guildId}),
 			(current) => {

fluxer_api/src/api/guild/repositories/GuildRepository.ts

@@ -12,6 +12,7 @@ import type {
 	GuildRow,
 	GuildStickerRow,
 } from '../../database/types/GuildTypes';
+import type {RequestCache} from '../../middleware/RequestCacheMiddleware';
 import type {Guild} from '../../models/Guild';
 import type {GuildAuditLog} from '../../models/GuildAuditLog';
 import type {GuildBan} from '../../models/GuildBan';
@@ -33,11 +34,11 @@ export class GuildRepository implements IGuildRepositoryAggregate {
 	private moderationRepo: GuildModerationRepository;
 	private contentRepo: GuildContentRepository;
 
-	constructor() {
-		this.dataRepo = new GuildDataRepository();
+	constructor(requestCache?: RequestCache) {
+		this.dataRepo = new GuildDataRepository(requestCache);
 		this.memberRepo = new GuildMemberRepository();
 		this.roleRepo = new GuildRoleRepository();
-		this.moderationRepo = new GuildModerationRepository();
+		this.moderationRepo = new GuildModerationRepository(requestCache);
 		this.contentRepo = new GuildContentRepository();
 	}
 

fluxer_api/src/api/middleware/GuildAvailabilityMiddleware.ts

@@ -64,6 +64,7 @@ async function resolveGuildIdForRequest(ctx: Context<HonoEnv>, path: string): Pr
 		return null;
 	}
 	const channel = await ctx.get('channelRepository').findUnique(channelId);
+	ctx.get('requestCache')?.channels.set(channelId, channel);
 	return channel?.guildId ?? null;
 }
 
@@ -81,6 +82,7 @@ export const GuildAvailabilityMiddleware = createMiddleware<HonoEnv>(async (ctx,
 	}
 	try {
 		const guild = await ctx.get('guildService').data.getGuildSystem(guildId);
+		ctx.get('requestCache')?.guilds.set(guildId, guild);
 		if (isGuildUnavailableForUser(guild, user)) {
 			throw new MissingAccessError();
 		}

fluxer_api/src/api/middleware/RequestCacheMiddleware.ts

@@ -2,21 +2,50 @@
 
 import type {UserPartialResponse} from '@fluxer/schema/src/domains/user/UserResponseSchemas';
 import {createMiddleware} from 'hono/factory';
+import type {ChannelID, GuildID} from '../BrandedTypes';
+import type {Channel} from '../models/Channel';
+import type {Guild} from '../models/Guild';
 import type {HonoEnv} from '../types/HonoEnv';
 
 export interface RequestCache {
 	userPartials: Map<bigint, UserPartialResponse>;
 	messageMentionChannels: Map<string, Array<{id: string; name: string; type: number}>>;
+	channels: Map<ChannelID, Channel | null>;
+	guilds: Map<GuildID, Guild>;
+	takeChannel(channelId: ChannelID): Channel | null | undefined;
+	takeGuild(guildId: GuildID): Guild | undefined;
 	clear(): void;
 }
 
 class RequestCacheImpl implements RequestCache {
 	userPartials = new Map<bigint, UserPartialResponse>();
 	messageMentionChannels = new Map<string, Array<{id: string; name: string; type: number}>>();
+	channels = new Map<ChannelID, Channel | null>();
+	guilds = new Map<GuildID, Guild>();
+
+	takeChannel(channelId: ChannelID): Channel | null | undefined {
+		if (!this.channels.has(channelId)) {
+			return undefined;
+		}
+		const channel = this.channels.get(channelId) ?? null;
+		this.channels.delete(channelId);
+		return channel;
+	}
+
+	takeGuild(guildId: GuildID): Guild | undefined {
+		const guild = this.guilds.get(guildId);
+		if (guild === undefined) {
+			return undefined;
+		}
+		this.guilds.delete(guildId);
+		return guild;
+	}
 
 	clear(): void {
 		this.userPartials.clear();
 		this.messageMentionChannels.clear();
+		this.channels.clear();
+		this.guilds.clear();
 	}
 }
 

fluxer_api/src/api/middleware/ServiceMiddleware.ts

@@ -17,6 +17,7 @@ import {SsoService} from '../auth/services/SsoService';
 import type {IBlueskyOAuthService} from '../bluesky/IBlueskyOAuthService';
 import {Config} from '../Config';
 import {createApiContext} from '../CreateApiContext';
+import {ChannelRepository} from '../channel/ChannelRepository';
 import {ChannelRequestService} from '../channel/services/ChannelRequestService';
 import {MessageRequestService} from '../channel/services/message/MessageRequestService';
 import {createMessageResponseDataService} from '../channel/services/message/MessageResponseDataService';
@@ -29,6 +30,7 @@ import {DonationCheckoutService} from '../donation/services/DonationCheckoutServ
 import {DonationMagicLinkService} from '../donation/services/DonationMagicLinkService';
 import {FavoriteMemeRequestService} from '../favorite_meme/FavoriteMemeRequestService';
 import {FavoriteMemeService} from '../favorite_meme/FavoriteMemeService';
+import {GuildRepository} from '../guild/repositories/GuildRepository';
 import {DisabledLiveKitService} from '../infrastructure/DisabledLiveKitService';
 import type {IGatewayService} from '../infrastructure/IGatewayService';
 import type {ILiveKitService} from '../infrastructure/ILiveKitService';
@@ -83,6 +85,7 @@ import {WebhookRequestService} from '../webhook/WebhookRequestService';
 import {WebhookService} from '../webhook/WebhookService';
 import {createGuildStackServices, type GuildStackServices} from './GuildStackServiceFactory';
 import {installLazyServices, type RequestScopedServices} from './LazyServiceProvider';
+import type {RequestCache} from './RequestCacheMiddleware';
 import {
 	ensureVoiceResourcesInitialized,
 	getBillingRepository,
@@ -348,6 +351,8 @@ class RequestServices implements RequestScopedServices {
 	private cachedVoiceRooms: IVoiceRoomStore | undefined;
 	private cachedVoice: VoiceService | null | undefined;
 	private cachedGuildStack: GuildStackServices | undefined;
+	private cachedChannelRepository: ChannelRepository | undefined;
+	private cachedGuildRepository: GuildRepository | undefined;
 	private cachedAdminService: AdminService | undefined;
 	private cachedApplicationService: ApplicationService | undefined;
 	private cachedAuthRequestService: AuthRequestService | undefined;
@@ -386,8 +391,14 @@ class RequestServices implements RequestScopedServices {
 	constructor(
 		private readonly context: ApiContext,
 		private readonly bluesky: IBlueskyOAuthService,
+		private readonly entityCache: RequestCache | undefined,
 	) {}
 
+	private get requestGuildRepository(): GuildRepository {
+		this.cachedGuildRepository ??= this.entityCache ? new GuildRepository(this.entityCache) : getGuildRepository();
+		return this.cachedGuildRepository;
+	}
+
 	get gatewayService(): IGatewayService {
 		this.cachedGateway ??= getGatewayService();
 		return this.cachedGateway;
@@ -437,9 +448,9 @@ class RequestServices implements RequestScopedServices {
 		this.cachedGuildStack ??= createGuildStackServices({
 			apiContext: this.context,
 			packRepository: getPackRepository(),
-			channelRepository: getChannelRepository(),
+			channelRepository: this.channelRepository,
 			userRepository: getUserRepository(),
-			guildRepository: getGuildRepository(),
+			guildRepository: this.requestGuildRepository,
 			inviteRepository: getInviteRepository(),
 			webhookRepository: getWebhookRepository(),
 			favoriteMemeRepository: getFavoriteMemeRepository(),
@@ -504,8 +515,11 @@ class RequestServices implements RequestScopedServices {
 		return getCacheService();
 	}
 
-	get channelRepository() {
-		return getChannelRepository();
+	get channelRepository(): ChannelRepository {
+		this.cachedChannelRepository ??= this.entityCache
+			? new ChannelRepository(this.entityCache)
+			: getChannelRepository();
+		return this.cachedChannelRepository;
 	}
 
 	get contactChangeLogService() {
@@ -1058,7 +1072,7 @@ export const ServiceMiddleware = createMiddleware<HonoEnv>(async (ctx, next) =>
 	await ensureVirusScanInitialized();
 	await ensureVoiceResourcesInitialized();
 	const blueskyOAuthService = await resolveBlueskyOAuthService(getInstanceConfigRepository());
-	installLazyServices(ctx, new RequestServices(apiContext, blueskyOAuthService));
+	installLazyServices(ctx, new RequestServices(apiContext, blueskyOAuthService, ctx.get('requestCache')));
 	await next();
 });