mirror of
https://github.com/fluxerapp/fluxer.git
synced 2026-09-02 21:04:06 +03:00
fix(kv): bound multi key fan out by pipelining per hash slot (#2313)
This commit is contained in:
@@ -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([]);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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
@@ -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>;
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user