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

This commit is contained in:
Hampus
2026-09-01 02:59:00 +02:00
committed by GitHub
parent 3b5b20c139
commit cf3af50464
14 changed files with 506 additions and 86 deletions
@@ -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<Array<string>>;
@@ -36,6 +18,7 @@ function createRecordingProvider(): {
setex: async (key: string) => {
commands.push([key]);
},
isClustered: () => true,
pipeline: () => {
const keys: Array<string> = [];
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([]);
});
});
+107
View File
@@ -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<Array<string>>;
peakInFlight: number;
}
function createRecordingProvider(clustered: boolean): RecordingProvider {
const recorder: RecordingProvider = {
client: {} as IKVProvider,
batches: [],
peakInFlight: 0,
};
let inFlight = 0;
const trackRoundTrip = async (keys: Array<string>): Promise<void> => {
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<string> = [];
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);
});
});
+20 -8
View File
@@ -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<void> {
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<number> {
@@ -104,5 +104,6 @@ export interface IKVProvider {
}>;
pipeline(): IKVPipeline;
multi(): IKVPipeline;
isClustered(): boolean;
health(): Promise<boolean>;
}
+32 -3
View File
@@ -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(),
@@ -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);
}
@@ -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']},
]);
});
});
@@ -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']);
});
});
@@ -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> {
@@ -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();
@@ -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');
});
});
@@ -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]);
@@ -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 {
@@ -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;