From cf3af5046482bb5280607de848548aeb7a76198f Mon Sep 17 00:00:00 2001 From: Hampus Date: Tue, 1 Sep 2026 02:59:00 +0200 Subject: [PATCH] fix(kv): bound multi key fan out by pipelining per hash slot (#2313) --- .../src/__tests__/CacheClusterSlots.test.ts | 25 +--- .../cache/src/__tests__/CacheFanout.test.ts | 107 +++++++++++++++++ .../cache/src/providers/KVCacheProvider.ts | 28 +++-- fluxer_api/pkgs/kv_client/src/IKVProvider.ts | 1 + fluxer_api/pkgs/kv_client/src/KVClient.ts | 35 +++++- fluxer_api/pkgs/kv_client/src/KVHashSlots.ts | 68 +++++++++++ .../__tests__/KVClientClusterSlots.test.ts | 56 ++++----- .../src/__tests__/KVClientFanout.test.ts | 112 ++++++++++++++++++ .../api/infrastructure/KVActivityTracker.ts | 16 ++- .../KVActivityTrackerClusterSlots.test.ts | 5 +- .../tests/KVActivityTrackerFanout.test.ts | 107 +++++++++++++++++ .../tests/VoiceRoomStoreClusterSlots.test.ts | 5 +- .../test/mocks/BatchRecordingKVProvider.ts | 22 +--- .../src/api/test/mocks/MockKVProvider.ts | 5 + 14 files changed, 506 insertions(+), 86 deletions(-) create mode 100644 fluxer_api/pkgs/cache/src/__tests__/CacheFanout.test.ts create mode 100644 fluxer_api/pkgs/kv_client/src/KVHashSlots.ts create mode 100644 fluxer_api/pkgs/kv_client/src/__tests__/KVClientFanout.test.ts create mode 100644 fluxer_api/src/api/infrastructure/tests/KVActivityTrackerFanout.test.ts diff --git a/fluxer_api/pkgs/cache/src/__tests__/CacheClusterSlots.test.ts b/fluxer_api/pkgs/cache/src/__tests__/CacheClusterSlots.test.ts index 62042abd5..4c5514de8 100644 --- a/fluxer_api/pkgs/cache/src/__tests__/CacheClusterSlots.test.ts +++ b/fluxer_api/pkgs/cache/src/__tests__/CacheClusterSlots.test.ts @@ -2,27 +2,9 @@ import {KVCacheProvider} from '@pkgs/cache/src/providers/KVCacheProvider'; import type {IKVPipeline, IKVProvider} from '@pkgs/kv_client/src/IKVProvider'; +import {computeHashSlot} from '@pkgs/kv_client/src/KVHashSlots'; import {describe, expect, it} from 'vitest'; -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 createRecordingProvider(): { client: IKVProvider; commands: Array>; @@ -36,6 +18,7 @@ function createRecordingProvider(): { setex: async (key: string) => { commands.push([key]); }, + isClustered: () => true, pipeline: () => { const keys: Array = []; commands.push(keys); @@ -61,7 +44,7 @@ describe('KVCacheProvider cluster hash slots', () => { const {client, commands} = createRecordingProvider(); const provider = new KVCacheProvider({client}); - expect(hashSlot('cache:alpha')).not.toBe(hashSlot('cache:beta')); + expect(computeHashSlot('cache:alpha')).not.toBe(computeHashSlot('cache:beta')); await provider.mset([ {key: 'cache:alpha', value: 1, ttlSeconds: 60}, @@ -69,6 +52,6 @@ describe('KVCacheProvider cluster hash slots', () => { ]); expect(commands.flat().sort()).toEqual(['cache:alpha', 'cache:beta']); - expect(commands.filter((keys) => new Set(keys.map(hashSlot)).size > 1)).toEqual([]); + expect(commands.filter((keys) => new Set(keys.map(computeHashSlot)).size > 1)).toEqual([]); }); }); diff --git a/fluxer_api/pkgs/cache/src/__tests__/CacheFanout.test.ts b/fluxer_api/pkgs/cache/src/__tests__/CacheFanout.test.ts new file mode 100644 index 000000000..c518bd292 --- /dev/null +++ b/fluxer_api/pkgs/cache/src/__tests__/CacheFanout.test.ts @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +import {KVCacheProvider} from '@pkgs/cache/src/providers/KVCacheProvider'; +import type {IKVPipeline, IKVProvider} from '@pkgs/kv_client/src/IKVProvider'; +import {computeHashSlot} from '@pkgs/kv_client/src/KVHashSlots'; +import {describe, expect, it} from 'vitest'; + +const MAX_CONCURRENT_ROUND_TRIPS = 16; + +interface RecordingProvider { + client: IKVProvider; + batches: Array>; + peakInFlight: number; +} + +function createRecordingProvider(clustered: boolean): RecordingProvider { + const recorder: RecordingProvider = { + client: {} as IKVProvider, + batches: [], + peakInFlight: 0, + }; + let inFlight = 0; + const trackRoundTrip = async (keys: Array): Promise => { + recorder.batches.push(keys); + inFlight += 1; + recorder.peakInFlight = Math.max(recorder.peakInFlight, inFlight); + await new Promise((resolve) => setTimeout(resolve, 0)); + inFlight -= 1; + }; + recorder.client = { + isClustered: () => clustered, + set: async (key: string) => { + await trackRoundTrip([key]); + return 'OK'; + }, + setex: async (key: string) => { + await trackRoundTrip([key]); + }, + pipeline: () => { + const keys: Array = []; + const batch = { + set: (key: string) => { + keys.push(key); + return batch; + }, + setex: (key: string) => { + keys.push(key); + return batch; + }, + exec: async () => { + await trackRoundTrip(keys); + return []; + }, + } as unknown as IKVPipeline; + return batch; + }, + } as unknown as IKVProvider; + return recorder; +} + +function createEntries(count: number): Array<{key: string; value: number; ttlSeconds: number}> { + return Array.from({length: count}, (_unused, index) => ({ + key: `cache:entry:${index}`, + value: index, + ttlSeconds: 60, + })); +} + +describe('KVCacheProvider multi entry write fan out', () => { + it('writes every entry in one round trip outside cluster mode', async () => { + const recorder = createRecordingProvider(false); + const provider = new KVCacheProvider({client: recorder.client}); + + await provider.mset(createEntries(1000)); + + expect(recorder.batches.map((keys) => keys.length)).toEqual([1000]); + expect(recorder.peakInFlight).toBe(1); + }); + + it('surfaces a failed command inside a batched write', async () => { + const client = { + isClustered: () => false, + pipeline: () => { + const batch = { + set: () => batch, + setex: () => batch, + exec: async () => [[new Error('write rejected'), null]], + } as unknown as IKVPipeline; + return batch; + }, + } as unknown as IKVProvider; + const provider = new KVCacheProvider({client}); + + await expect(provider.mset(createEntries(2))).rejects.toThrow('write rejected'); + }); + + it('bounds concurrent round trips when entries span hash slots', async () => { + const recorder = createRecordingProvider(true); + const provider = new KVCacheProvider({client: recorder.client}); + + await provider.mset(createEntries(1000)); + + expect(recorder.peakInFlight).toBeLessThanOrEqual(MAX_CONCURRENT_ROUND_TRIPS); + expect(recorder.batches.filter((keys) => new Set(keys.map(computeHashSlot)).size > 1)).toEqual([]); + expect(recorder.batches.flat().length).toBe(1000); + }); +}); diff --git a/fluxer_api/pkgs/cache/src/providers/KVCacheProvider.ts b/fluxer_api/pkgs/cache/src/providers/KVCacheProvider.ts index 444c5e880..60316a064 100644 --- a/fluxer_api/pkgs/cache/src/providers/KVCacheProvider.ts +++ b/fluxer_api/pkgs/cache/src/providers/KVCacheProvider.ts @@ -11,6 +11,7 @@ import type {CacheLogger, CacheTelemetry} from '@pkgs/cache/src/CacheProviderTyp import {parseCachedValue, safeJsonParse, serializeValue} from '@pkgs/cache/src/CacheSerialization'; import {type CacheLookupResult, ICacheService} from '@pkgs/cache/src/ICacheService'; import type {IKVProvider} from '@pkgs/kv_client/src/IKVProvider'; +import {runSlotBatches, splitIntoSlotBatches} from '@pkgs/kv_client/src/KVHashSlots'; interface KVCacheProviderConfig { client: IKVProvider; @@ -137,16 +138,27 @@ export class KVCacheProvider extends ICacheService { }>, ): Promise { if (entries.length === 0) return; - await Promise.all( - entries.map(async (entry) => { - const serialized = serializeValue(entry.value); + const serialized = entries.map((entry) => ({ + key: entry.key, + value: serializeValue(entry.value), + ttlSeconds: entry.ttlSeconds, + })); + const batches = splitIntoSlotBatches(serialized, (entry) => entry.key, this.client.isClustered()); + await runSlotBatches(batches, async (batch) => { + const pipeline = this.client.pipeline(); + for (const entry of batch) { if (entry.ttlSeconds) { - await this.client.setex(entry.key, entry.ttlSeconds, serialized); - return; + pipeline.setex(entry.key, entry.ttlSeconds, entry.value); + } else { + pipeline.set(entry.key, entry.value); } - await this.client.set(entry.key, serialized); - }), - ); + } + for (const [error] of await pipeline.exec()) { + if (error) { + throw error; + } + } + }); } async deletePattern(pattern: string): Promise { diff --git a/fluxer_api/pkgs/kv_client/src/IKVProvider.ts b/fluxer_api/pkgs/kv_client/src/IKVProvider.ts index 130b5c993..988e235a4 100644 --- a/fluxer_api/pkgs/kv_client/src/IKVProvider.ts +++ b/fluxer_api/pkgs/kv_client/src/IKVProvider.ts @@ -104,5 +104,6 @@ export interface IKVProvider { }>; pipeline(): IKVPipeline; multi(): IKVPipeline; + isClustered(): boolean; health(): Promise; } diff --git a/fluxer_api/pkgs/kv_client/src/KVClient.ts b/fluxer_api/pkgs/kv_client/src/KVClient.ts index 1fca0be6e..8110aafaa 100644 --- a/fluxer_api/pkgs/kv_client/src/KVClient.ts +++ b/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(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): Promise { @@ -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 = []; + 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(items: ReadonlyArray, keyOf: (item: T) => string): Array> { + return splitIntoSlotBatches(items, keyOf, this.isClustered()); + } + pipeline(): IKVPipeline { return new KVPipeline({ createCommander: () => this.client.pipeline(), diff --git a/fluxer_api/pkgs/kv_client/src/KVHashSlots.ts b/fluxer_api/pkgs/kv_client/src/KVHashSlots.ts new file mode 100644 index 000000000..bfcf6be3b --- /dev/null +++ b/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( + items: ReadonlyArray, + keyOf: (item: T) => string, + clustered: boolean, +): Array> { + if (items.length === 0) { + return []; + } + if (!clustered) { + return [[...items]]; + } + const batches = new Map>(); + 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(batches: ReadonlyArray, run: (batch: T) => Promise): Promise { + 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); +} diff --git a/fluxer_api/pkgs/kv_client/src/__tests__/KVClientClusterSlots.test.ts b/fluxer_api/pkgs/kv_client/src/__tests__/KVClientClusterSlots.test.ts index cb5093ebc..a79601bea 100644 --- a/fluxer_api/pkgs/kv_client/src/__tests__/KVClientClusterSlots.test.ts +++ b/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}> { - 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']}, + ]); + }); }); diff --git a/fluxer_api/pkgs/kv_client/src/__tests__/KVClientFanout.test.ts b/fluxer_api/pkgs/kv_client/src/__tests__/KVClientFanout.test.ts new file mode 100644 index 000000000..03801ce27 --- /dev/null +++ b/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}>, + store: new Map(), + tracker: {inFlight: 0, peakInFlight: 0}, +})); + +vi.mock('ioredis', () => { + const trackRoundTrip = async (name: string, keys: Array): Promise => { + 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): Promise> { + await trackRoundTrip('mget', keys); + return keys.map((key) => store.get(key) ?? null); + } + + async mset(...args: Array): Promise { + const keys: Array = []; + 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): Promise { + await trackRoundTrip('del', keys); + return keys.length; + } + + async get(key: string): Promise { + await trackRoundTrip('get', [key]); + return store.get(key) ?? null; + } + + async set(key: string, value: string): Promise { + await trackRoundTrip('set', [key]); + store.set(key, value); + return 'OK'; + } + } + return {default: MockRedis, Cluster: MockRedis}; +}); + +function createKeys(count: number): Array { + return Array.from({length: count}, (_unused, index) => `fanout:key:${index}`); +} + +function crossSlotCommands(): Array<{name: string; keys: Array}> { + 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']); + }); +}); diff --git a/fluxer_api/src/api/infrastructure/KVActivityTracker.ts b/fluxer_api/src/api/infrastructure/KVActivityTracker.ts index 36209b205..604de3304 100644 --- a/fluxer_api/src/api/infrastructure/KVActivityTracker.ts +++ b/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 { - 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 { + 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 { diff --git a/fluxer_api/src/api/infrastructure/tests/KVActivityTrackerClusterSlots.test.ts b/fluxer_api/src/api/infrastructure/tests/KVActivityTrackerClusterSlots.test.ts index aca0b4343..964ed168c 100644 --- a/fluxer_api/src/api/infrastructure/tests/KVActivityTrackerClusterSlots.test.ts +++ b/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(); diff --git a/fluxer_api/src/api/infrastructure/tests/KVActivityTrackerFanout.test.ts b/fluxer_api/src/api/infrastructure/tests/KVActivityTrackerFanout.test.ts new file mode 100644 index 000000000..8fec036e1 --- /dev/null +++ b/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 { + 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(run: () => Promise): Promise { + 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 { + 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'); + }); +}); diff --git a/fluxer_api/src/api/infrastructure/tests/VoiceRoomStoreClusterSlots.test.ts b/fluxer_api/src/api/infrastructure/tests/VoiceRoomStoreClusterSlots.test.ts index b5debb4b3..ef626c856 100644 --- a/fluxer_api/src/api/infrastructure/tests/VoiceRoomStoreClusterSlots.test.ts +++ b/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]); diff --git a/fluxer_api/src/api/test/mocks/BatchRecordingKVProvider.ts b/fluxer_api/src/api/test/mocks/BatchRecordingKVProvider.ts index 477b1322b..d7b0b937a 100644 --- a/fluxer_api/src/api/test/mocks/BatchRecordingKVProvider.ts +++ b/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; } -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 = []; @@ -41,7 +23,7 @@ export class BatchRecordingKVProvider extends MockKVProvider { } crossSlotBatches(): Array { - 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 { diff --git a/fluxer_api/src/api/test/mocks/MockKVProvider.ts b/fluxer_api/src/api/test/mocks/MockKVProvider.ts index d1ee83c0a..64f5e22c9 100644 --- a/fluxer_api/src/api/test/mocks/MockKVProvider.ts +++ b/fluxer_api/src/api/test/mocks/MockKVProvider.ts @@ -111,6 +111,7 @@ export class MockKVProvider implements IKVProvider { key: string; values: Array; }> = []; + clustered = true; private subscription: MockKVSubscription; private readonly stringStore = new Map(); private readonly setStore = new Map>(); @@ -782,6 +783,10 @@ export class MockKVProvider implements IKVProvider { return this.createPipeline(); } + isClustered(): boolean { + return this.clustered; + } + async health(): Promise { this.healthSpy(); return true;