fix(kv): bound multi key fan out by pipelining per hash slot (#2313)

fluxer_api/src/api/infrastructure/KVActivityTracker.ts

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
 import type {IKVProvider} from '@pkgs/kv_client/src/IKVProvider';
+import {runSlotBatches, splitIntoSlotBatches} from '@pkgs/kv_client/src/KVHashSlots';
 import {seconds} from 'itty-time';
 import type {UserID} from '../BrandedTypes';
 import {Logger} from '../Logger';
@@ -74,8 +75,19 @@ export class KVActivityTracker {
 		return age > STATE_VERSION_TTL_SECONDS;
 	}
 
-	private async writeActivityBatch(batch: ReadonlyArray<{key: string; value: string}>): Promise<void> {
-		await Promise.all(batch.map(async (entry) => this.kvClient.setex(entry.key, TTL_SECONDS, entry.value)));
+	private async writeActivityBatch(entries: ReadonlyArray<{key: string; value: string}>): Promise<void> {
+		const batches = splitIntoSlotBatches(entries, (entry) => entry.key, this.kvClient.isClustered());
+		await runSlotBatches(batches, async (batch) => {
+			const pipeline = this.kvClient.pipeline();
+			for (const entry of batch) {
+				pipeline.setex(entry.key, TTL_SECONDS, entry.value);
+			}
+			for (const [error] of await pipeline.exec()) {
+				if (error) {
+					throw error;
+				}
+			}
+		});
 	}
 
 	async rebuildActivities(): Promise<void> {

fluxer_api/src/api/infrastructure/tests/KVActivityTrackerClusterSlots.test.ts

@@ -1,9 +1,10 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
+import {computeHashSlot} from '@pkgs/kv_client/src/KVHashSlots';
 import {afterEach, describe, expect, it, vi} from 'vitest';
 import {createUserID} from '../../BrandedTypes';
 import type {User} from '../../models/User';
-import {BatchRecordingKVProvider, hashSlot} from '../../test/mocks/BatchRecordingKVProvider';
+import {BatchRecordingKVProvider} from '../../test/mocks/BatchRecordingKVProvider';
 import {UserRepository} from '../../user/repositories/UserRepository';
 import {KVActivityTracker} from '../KVActivityTracker';
 
@@ -22,7 +23,7 @@ describe('KVActivityTracker cluster hash slots', () => {
 		const users = [createUser(1234n, lastActiveAt), createUser(5678n, lastActiveAt)];
 		vi.spyOn(UserRepository.prototype, 'scanAllUsersPage').mockResolvedValue({users, pageState: null});
 
-		expect(hashSlot('user_activity:1234')).not.toBe(hashSlot('user_activity:5678'));
+		expect(computeHashSlot('user_activity:1234')).not.toBe(computeHashSlot('user_activity:5678'));
 
 		await new KVActivityTracker(kvClient).rebuildActivities();
 

fluxer_api/src/api/infrastructure/tests/KVActivityTrackerFanout.test.ts

@@ -0,0 +1,107 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import type {IKVPipeline} from '@pkgs/kv_client/src/IKVProvider';
+import {afterEach, describe, expect, it, vi} from 'vitest';
+import {createUserID} from '../../BrandedTypes';
+import type {User} from '../../models/User';
+import {BatchRecordingKVProvider} from '../../test/mocks/BatchRecordingKVProvider';
+import {UserRepository} from '../../user/repositories/UserRepository';
+import {KVActivityTracker} from '../KVActivityTracker';
+
+const REBUILD_KV_BATCH_SIZE = 1000;
+const MAX_CONCURRENT_ROUND_TRIPS = 16;
+
+class FanoutRecordingKVProvider extends BatchRecordingKVProvider {
+	inFlight = 0;
+	peakInFlight = 0;
+	failingKey: string | null = null;
+
+	private insidePipeline = 0;
+
+	override pipeline(): IKVPipeline {
+		const inner = super.pipeline();
+		return {
+			...inner,
+			exec: async () =>
+				await this.trackRoundTrip(async () => {
+					this.insidePipeline += 1;
+					try {
+						return await inner.exec();
+					} finally {
+						this.insidePipeline -= 1;
+					}
+				}),
+		};
+	}
+
+	override async setex(key: string, ttlSeconds: number, value: string): Promise<void> {
+		if (key === this.failingKey) {
+			throw new Error('write rejected');
+		}
+		if (this.insidePipeline > 0) {
+			await super.setex(key, ttlSeconds, value);
+			return;
+		}
+		await this.trackRoundTrip(async () => {
+			await super.setex(key, ttlSeconds, value);
+		});
+	}
+
+	private async trackRoundTrip<T>(run: () => Promise<T>): Promise<T> {
+		this.inFlight += 1;
+		this.peakInFlight = Math.max(this.peakInFlight, this.inFlight);
+		try {
+			await new Promise((resolve) => setTimeout(resolve, 0));
+			return await run();
+		} finally {
+			this.inFlight -= 1;
+		}
+	}
+}
+
+function mockUserPage(count: number): Array<User> {
+	const lastActiveAt = new Date('2026-06-01T00:00:00.000Z');
+	const users = Array.from(
+		{length: count},
+		(_unused, index) => ({id: createUserID(BigInt(index + 1)), lastActiveAt}) as unknown as User,
+	);
+	vi.spyOn(UserRepository.prototype, 'scanAllUsersPage').mockResolvedValue({users, pageState: null});
+	return users;
+}
+
+describe('KVActivityTracker rebuild fan out', () => {
+	afterEach(() => {
+		vi.restoreAllMocks();
+	});
+
+	it('writes a full rebuild batch in one round trip outside cluster mode', async () => {
+		const kvClient = new FanoutRecordingKVProvider();
+		kvClient.clustered = false;
+		mockUserPage(REBUILD_KV_BATCH_SIZE);
+
+		await new KVActivityTracker(kvClient).rebuildActivities();
+
+		expect(kvClient.batches.map((batch) => batch.keys.length)).toEqual([REBUILD_KV_BATCH_SIZE]);
+		expect(kvClient.peakInFlight).toBe(1);
+	});
+
+	it('bounds concurrent round trips when a rebuild batch spans hash slots', async () => {
+		const kvClient = new FanoutRecordingKVProvider();
+		const users = mockUserPage(REBUILD_KV_BATCH_SIZE);
+
+		await new KVActivityTracker(kvClient).rebuildActivities();
+
+		expect(kvClient.peakInFlight).toBeLessThanOrEqual(MAX_CONCURRENT_ROUND_TRIPS);
+		expect(kvClient.crossSlotBatches()).toEqual([]);
+		expect(kvClient.batches.flatMap((batch) => batch.keys).length).toBe(REBUILD_KV_BATCH_SIZE);
+		expect(await kvClient.get(`user_activity:${users[0].id}`)).toBe(users[0].lastActiveAt?.getTime().toString());
+	});
+
+	it('fails the rebuild when a batched write fails', async () => {
+		const kvClient = new FanoutRecordingKVProvider();
+		const users = mockUserPage(REBUILD_KV_BATCH_SIZE);
+		kvClient.failingKey = `user_activity:${users[0].id}`;
+
+		await expect(new KVActivityTracker(kvClient).rebuildActivities()).rejects.toThrow('write rejected');
+	});
+});

fluxer_api/src/api/infrastructure/tests/VoiceRoomStoreClusterSlots.test.ts

@@ -1,8 +1,9 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
+import {computeHashSlot} from '@pkgs/kv_client/src/KVHashSlots';
 import {describe, expect, it} from 'vitest';
 import {createChannelID, createGuildID} from '../../BrandedTypes';
-import {BatchRecordingKVProvider, hashSlot} from '../../test/mocks/BatchRecordingKVProvider';
+import {BatchRecordingKVProvider} from '../../test/mocks/BatchRecordingKVProvider';
 import {VOICE_OCCUPANCY_REGION_KEY_PREFIX, VOICE_OCCUPANCY_SERVER_KEY_PREFIX} from '../../voice/VoiceConstants';
 import {VoiceRoomStore} from '../VoiceRoomStore';
 
@@ -16,7 +17,7 @@ describe('VoiceRoomStore cluster hash slots', () => {
 		const serverKey = `${VOICE_OCCUPANCY_SERVER_KEY_PREFIX}:us-east:voice-1`;
 		const member = 'guild:1234:channel:5678';
 
-		expect(hashSlot(regionKey)).not.toBe(hashSlot(serverKey));
+		expect(computeHashSlot(regionKey)).not.toBe(computeHashSlot(serverKey));
 
 		await store.pinRoomServer(guildId, channelId, 'us-east', 'voice-1', 'wss://voice-1.example');
 		expect(await kvClient.smembers(regionKey)).toEqual([member]);

fluxer_api/src/api/test/mocks/BatchRecordingKVProvider.ts

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
 import type {IKVPipeline} from '@pkgs/kv_client/src/IKVProvider';
+import {computeHashSlot} from '@pkgs/kv_client/src/KVHashSlots';
 import {MockKVProvider} from './MockKVProvider';
 
 type BatchMode = 'multi' | 'pipeline';
@@ -10,25 +11,6 @@ interface RecordedBatch {
 	keys: Array<string>;
 }
 
-export function hashSlot(key: string): number {
-	let hashed = key;
-	const start = key.indexOf('{');
-	if (start !== -1) {
-		const end = key.indexOf('}', start + 1);
-		if (end > start + 1) {
-			hashed = key.slice(start + 1, end);
-		}
-	}
-	let crc = 0;
-	for (let index = 0; index < hashed.length; index += 1) {
-		crc ^= (hashed.charCodeAt(index) & 0xff) << 8;
-		for (let bit = 0; bit < 8; bit += 1) {
-			crc = (crc & 0x8000) === 0 ? (crc << 1) & 0xffff : ((crc << 1) ^ 0x1021) & 0xffff;
-		}
-	}
-	return crc % 16384;
-}
-
 export class BatchRecordingKVProvider extends MockKVProvider {
 	readonly batches: Array<RecordedBatch> = [];
 
@@ -41,7 +23,7 @@ export class BatchRecordingKVProvider extends MockKVProvider {
 	}
 
 	crossSlotBatches(): Array<RecordedBatch> {
-		return this.batches.filter((batch) => new Set(batch.keys.map(hashSlot)).size > 1);
+		return this.batches.filter((batch) => new Set(batch.keys.map(computeHashSlot)).size > 1);
 	}
 
 	private recordBatch(mode: BatchMode, inner: IKVPipeline): IKVPipeline {

fluxer_api/src/api/test/mocks/MockKVProvider.ts

@@ -111,6 +111,7 @@ export class MockKVProvider implements IKVProvider {
 		key: string;
 		values: Array<string>;
 	}> = [];
+	clustered = true;
 	private subscription: MockKVSubscription;
 	private readonly stringStore = new Map<string, string>();
 	private readonly setStore = new Map<string, Set<string>>();
@@ -782,6 +783,10 @@ export class MockKVProvider implements IKVProvider {
 		return this.createPipeline();
 	}
 
+	isClustered(): boolean {
+		return this.clustered;
+	}
+
 	async health(): Promise<boolean> {
 		this.healthSpy();
 		return true;