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

fluxer_api/pkgs/kv_client/src/IKVProvider.ts

@@ -104,5 +104,6 @@ export interface IKVProvider {
 	}>;
 	pipeline(): IKVPipeline;
 	multi(): IKVPipeline;
+	isClustered(): boolean;
 	health(): Promise<boolean>;
 }

fluxer_api/pkgs/kv_client/src/KVClient.ts

@@ -16,6 +16,7 @@ import {
 	parseRangeByScoreArguments,
 	parseSetArguments,
 } from '@pkgs/kv_client/src/KVCommandArguments';
+import {runSlotBatches, splitIntoSlotBatches} from '@pkgs/kv_client/src/KVHashSlots';
 import {KVPipeline} from '@pkgs/kv_client/src/KVPipeline';
 import {KVSubscription} from '@pkgs/kv_client/src/KVSubscription';
 import Redis, {Cluster} from 'ioredis';
@@ -358,7 +359,20 @@ export class KVClient implements IKVProvider {
 		if (keys.length === 0) {
 			return [];
 		}
-		return await this.execute('mget', async () => await Promise.all(keys.map(async (key) => this.client.get(key))));
+		return await this.execute('mget', async () => {
+			const values = new Array<string | null>(keys.length).fill(null);
+			const batches = this.splitBySlot(
+				keys.map((key, index) => ({key, index})),
+				(entry) => entry.key,
+			);
+			await runSlotBatches(batches, async (batch) => {
+				const batchValues = await this.client.mget(...batch.map((entry) => entry.key));
+				for (const [position, entry] of batch.entries()) {
+					values[entry.index] = batchValues[position] ?? null;
+				}
+			});
+			return values;
+		});
 	}
 
 	async mset(...args: Array<string>): Promise<void> {
@@ -367,7 +381,10 @@ export class KVClient implements IKVProvider {
 			return;
 		}
 		await this.execute('mset', async () => {
-			await Promise.all(entries.map(async (entry) => this.client.set(entry.key, entry.value)));
+			const batches = this.splitBySlot(entries, (entry) => entry.key);
+			await runSlotBatches(batches, async (batch) => {
+				await this.client.mset(...batch.flatMap((entry) => [entry.key, entry.value]));
+			});
 		});
 	}
 
@@ -376,7 +393,11 @@ export class KVClient implements IKVProvider {
 			return 0;
 		}
 		return await this.execute('del', async () => {
-			const deleted = await Promise.all(keys.map(async (key) => this.client.del(key)));
+			const deleted: Array<number> = [];
+			const batches = this.splitBySlot(keys, (key) => key);
+			await runSlotBatches(batches, async (batch) => {
+				deleted.push(await this.client.del(...batch));
+			});
 			return deleted.reduce((total, count) => total + count, 0);
 		});
 	}
@@ -707,6 +728,14 @@ export class KVClient implements IKVProvider {
 		});
 	}
 
+	isClustered(): boolean {
+		return this.config.mode === 'cluster';
+	}
+
+	private splitBySlot<T>(items: ReadonlyArray<T>, keyOf: (item: T) => string): Array<Array<T>> {
+		return splitIntoSlotBatches(items, keyOf, this.isClustered());
+	}
+
 	pipeline(): IKVPipeline {
 		return new KVPipeline({
 			createCommander: () => this.client.pipeline(),

fluxer_api/pkgs/kv_client/src/KVHashSlots.ts

@@ -0,0 +1,68 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+const HASH_SLOT_COUNT = 16384;
+const MAX_CONCURRENT_SLOT_BATCHES = 16;
+
+function extractHashTag(key: string): string {
+	const start = key.indexOf('{');
+	if (start === -1) {
+		return key;
+	}
+	const end = key.indexOf('}', start + 1);
+	if (end > start + 1) {
+		return key.slice(start + 1, end);
+	}
+	return key;
+}
+
+export function computeHashSlot(key: string): number {
+	const hashed = extractHashTag(key);
+	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 % HASH_SLOT_COUNT;
+}
+
+export function splitIntoSlotBatches<T>(
+	items: ReadonlyArray<T>,
+	keyOf: (item: T) => string,
+	clustered: boolean,
+): Array<Array<T>> {
+	if (items.length === 0) {
+		return [];
+	}
+	if (!clustered) {
+		return [[...items]];
+	}
+	const batches = new Map<number, Array<T>>();
+	for (const item of items) {
+		const slot = computeHashSlot(keyOf(item));
+		const batch = batches.get(slot);
+		if (batch) {
+			batch.push(item);
+		} else {
+			batches.set(slot, [item]);
+		}
+	}
+	return [...batches.values()];
+}
+
+export async function runSlotBatches<T>(batches: ReadonlyArray<T>, run: (batch: T) => Promise<void>): Promise<void> {
+	if (batches.length <= MAX_CONCURRENT_SLOT_BATCHES) {
+		await Promise.all(batches.map(async (batch) => await run(batch)));
+		return;
+	}
+	let nextIndex = 0;
+	const workers = Array.from({length: MAX_CONCURRENT_SLOT_BATCHES}, async () => {
+		while (nextIndex < batches.length) {
+			const batch = batches[nextIndex];
+			nextIndex += 1;
+			await run(batch);
+		}
+	});
+	await Promise.all(workers);
+}

fluxer_api/pkgs/kv_client/src/__tests__/KVClientClusterSlots.test.ts

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
 import {KVClient} from '@pkgs/kv_client/src/KVClient';
+import {computeHashSlot} from '@pkgs/kv_client/src/KVHashSlots';
 import {beforeEach, describe, expect, it, vi} from 'vitest';
 
 const {commands, store} = vi.hoisted(() => ({
@@ -44,31 +45,16 @@ vi.mock('ioredis', () => {
 	return {default: MockRedis, Cluster: MockRedis};
 });
 
-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;
-}
-
 function crossSlotCommands(): Array<{name: string; keys: Array<string>}> {
-	return commands.filter((command) => new Set(command.keys.map(hashSlot)).size > 1);
+	return commands.filter((command) => new Set(command.keys.map(computeHashSlot)).size > 1);
 }
 
-function createClient(): KVClient {
-	return new KVClient('redis://127.0.0.1:6379');
+function createClusteredClient(): KVClient {
+	return new KVClient({url: 'redis://127.0.0.1:6379', mode: 'cluster'});
+}
+
+function createStandaloneClient(): KVClient {
+	return new KVClient({url: 'redis://127.0.0.1:6379', mode: 'standalone'});
 }
 
 describe('KVClient cluster hash slots', () => {
@@ -78,8 +64,8 @@ describe('KVClient cluster hash slots', () => {
 	});
 
 	it('reads several keys without a command spanning hash slots', async () => {
-		expect(hashSlot('slot:alpha')).not.toBe(hashSlot('slot:beta'));
-		const client = createClient();
+		expect(computeHashSlot('slot:alpha')).not.toBe(computeHashSlot('slot:beta'));
+		const client = createClusteredClient();
 		await client.set('slot:alpha', 'one');
 
 		await expect(client.mget('slot:alpha', 'slot:beta')).resolves.toEqual(['one', null]);
@@ -87,8 +73,8 @@ describe('KVClient cluster hash slots', () => {
 	});
 
 	it('writes several keys without a command spanning hash slots', async () => {
-		expect(hashSlot('slot:alpha')).not.toBe(hashSlot('slot:beta'));
-		const client = createClient();
+		expect(computeHashSlot('slot:alpha')).not.toBe(computeHashSlot('slot:beta'));
+		const client = createClusteredClient();
 
 		await client.mset('slot:alpha', 'one', 'slot:beta', 'two');
 
@@ -98,8 +84,8 @@ describe('KVClient cluster hash slots', () => {
 	});
 
 	it('deletes several keys without a command spanning hash slots', async () => {
-		expect(hashSlot('slot:alpha')).not.toBe(hashSlot('slot:beta'));
-		const client = createClient();
+		expect(computeHashSlot('slot:alpha')).not.toBe(computeHashSlot('slot:beta'));
+		const client = createClusteredClient();
 		await client.set('slot:alpha', 'one');
 		await client.set('slot:beta', 'two');
 
@@ -107,4 +93,18 @@ describe('KVClient cluster hash slots', () => {
 		expect(store.size).toBe(0);
 		expect(crossSlotCommands()).toEqual([]);
 	});
+
+	it('keeps multi key commands whole outside cluster mode', async () => {
+		const client = createStandaloneClient();
+
+		await client.mset('slot:alpha', 'one', 'slot:beta', 'two');
+		await expect(client.mget('slot:alpha', 'slot:beta')).resolves.toEqual(['one', 'two']);
+		await expect(client.del('slot:alpha', 'slot:beta')).resolves.toBe(2);
+
+		expect(commands).toEqual([
+			{name: 'mset', keys: ['slot:alpha', 'slot:beta']},
+			{name: 'mget', keys: ['slot:alpha', 'slot:beta']},
+			{name: 'del', keys: ['slot:alpha', 'slot:beta']},
+		]);
+	});
 });

fluxer_api/pkgs/kv_client/src/__tests__/KVClientFanout.test.ts

@@ -0,0 +1,112 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import {KVClient} from '@pkgs/kv_client/src/KVClient';
+import {computeHashSlot} from '@pkgs/kv_client/src/KVHashSlots';
+import {beforeEach, describe, expect, it, vi} from 'vitest';
+
+const MAX_CONCURRENT_ROUND_TRIPS = 16;
+
+const {commands, store, tracker} = vi.hoisted(() => ({
+	commands: [] as Array<{name: string; keys: Array<string>}>,
+	store: new Map<string, string>(),
+	tracker: {inFlight: 0, peakInFlight: 0},
+}));
+
+vi.mock('ioredis', () => {
+	const trackRoundTrip = async (name: string, keys: Array<string>): Promise<void> => {
+		commands.push({name, keys});
+		tracker.inFlight += 1;
+		tracker.peakInFlight = Math.max(tracker.peakInFlight, tracker.inFlight);
+		await new Promise((resolve) => setTimeout(resolve, 0));
+		tracker.inFlight -= 1;
+	};
+	class MockRedis {
+		async mget(...keys: Array<string>): Promise<Array<string | null>> {
+			await trackRoundTrip('mget', keys);
+			return keys.map((key) => store.get(key) ?? null);
+		}
+
+		async mset(...args: Array<string>): Promise<string> {
+			const keys: Array<string> = [];
+			for (let index = 0; index + 1 < args.length; index += 2) {
+				keys.push(args[index]);
+				store.set(args[index], args[index + 1]);
+			}
+			await trackRoundTrip('mset', keys);
+			return 'OK';
+		}
+
+		async del(...keys: Array<string>): Promise<number> {
+			await trackRoundTrip('del', keys);
+			return keys.length;
+		}
+
+		async get(key: string): Promise<string | null> {
+			await trackRoundTrip('get', [key]);
+			return store.get(key) ?? null;
+		}
+
+		async set(key: string, value: string): Promise<string> {
+			await trackRoundTrip('set', [key]);
+			store.set(key, value);
+			return 'OK';
+		}
+	}
+	return {default: MockRedis, Cluster: MockRedis};
+});
+
+function createKeys(count: number): Array<string> {
+	return Array.from({length: count}, (_unused, index) => `fanout:key:${index}`);
+}
+
+function crossSlotCommands(): Array<{name: string; keys: Array<string>}> {
+	return commands.filter((command) => new Set(command.keys.map(computeHashSlot)).size > 1);
+}
+
+describe('KVClient multi key fan out', () => {
+	beforeEach(() => {
+		commands.length = 0;
+		store.clear();
+		tracker.inFlight = 0;
+		tracker.peakInFlight = 0;
+	});
+
+	it('bounds concurrent round trips for a cluster read', async () => {
+		const client = new KVClient({url: 'redis://127.0.0.1:6379', mode: 'cluster'});
+
+		const values = await client.mget(...createKeys(1000));
+
+		expect(values.length).toBe(1000);
+		expect(tracker.peakInFlight).toBeLessThanOrEqual(MAX_CONCURRENT_ROUND_TRIPS);
+		expect(crossSlotCommands()).toEqual([]);
+	});
+
+	it('bounds concurrent round trips for a cluster write', async () => {
+		const client = new KVClient({url: 'redis://127.0.0.1:6379', mode: 'cluster'});
+
+		await client.mset(...createKeys(1000).flatMap((key) => [key, 'value']));
+
+		expect(tracker.peakInFlight).toBeLessThanOrEqual(MAX_CONCURRENT_ROUND_TRIPS);
+		expect(crossSlotCommands()).toEqual([]);
+	});
+
+	it('reads a thousand keys in one round trip outside cluster mode', async () => {
+		const client = new KVClient({url: 'redis://127.0.0.1:6379', mode: 'standalone'});
+
+		await client.mget(...createKeys(1000));
+
+		expect(commands.map((command) => ({name: command.name, count: command.keys.length}))).toEqual([
+			{name: 'mget', count: 1000},
+		]);
+		expect(tracker.peakInFlight).toBe(1);
+	});
+
+	it('keeps values ordered when a cluster read is split by slot', async () => {
+		const client = new KVClient({url: 'redis://127.0.0.1:6379', mode: 'cluster'});
+		await client.mset('fanout:a', 'one', 'fanout:b', 'two');
+
+		const values = await client.mget('fanout:a', 'fanout:missing', 'fanout:b');
+
+		expect(values).toEqual(['one', null, 'two']);
+	});
+});