mirror of
https://github.com/discordjs/discord.js.git
synced 2026-03-18 04:23:31 +01:00
refactor: abstract identify throttling and correct max_concurrency handling (#9375)
* refactor: properly support max_concurrency ratelimit keys * fix: properly block for same key * chore: export session state * chore: throttler no longer requires manager * refactor: abstract throttlers * chore: proper member order * chore: remove leftover debug log * chore: use @link tag in doc comment Co-authored-by: Jiralite <33201955+Jiralite@users.noreply.github.com> * chore: suggested changes * fix(WebSocketShard): cancel identify if the shard closed in the meantime * refactor(throttlers): support abort signals * fix: memory leak * chore: remove leftover --------- Co-authored-by: Jiralite <33201955+Jiralite@users.noreply.github.com> Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>
This commit is contained in:
@@ -1,46 +0,0 @@
|
||||
import { setTimeout as sleep } from 'node:timers/promises';
|
||||
import { expect, test, vi, type Mock } from 'vitest';
|
||||
import { IdentifyThrottler, type WebSocketManager } from '../../src/index.js';
|
||||
|
||||
vi.mock('node:timers/promises', () => ({
|
||||
setTimeout: vi.fn(),
|
||||
}));
|
||||
|
||||
const fetchGatewayInformation = vi.fn();
|
||||
|
||||
const manager = {
|
||||
fetchGatewayInformation,
|
||||
} as unknown as WebSocketManager;
|
||||
|
||||
const throttler = new IdentifyThrottler(manager);
|
||||
|
||||
vi.useFakeTimers();
|
||||
|
||||
const NOW = vi.fn().mockReturnValue(Date.now());
|
||||
global.Date.now = NOW;
|
||||
|
||||
test('wait for identify', async () => {
|
||||
fetchGatewayInformation.mockReturnValue({
|
||||
session_start_limit: {
|
||||
max_concurrency: 2,
|
||||
},
|
||||
});
|
||||
|
||||
// First call should never wait
|
||||
await throttler.waitForIdentify();
|
||||
expect(sleep).not.toHaveBeenCalled();
|
||||
|
||||
// Second call still won't wait because max_concurrency is 2
|
||||
await throttler.waitForIdentify();
|
||||
expect(sleep).not.toHaveBeenCalled();
|
||||
|
||||
// Third call should wait
|
||||
await throttler.waitForIdentify();
|
||||
expect(sleep).toHaveBeenCalled();
|
||||
|
||||
(sleep as Mock).mockRestore();
|
||||
|
||||
// Fourth call shouldn't wait, because our max_concurrency is 2 and we waited for a reset
|
||||
await throttler.waitForIdentify();
|
||||
expect(sleep).not.toHaveBeenCalled();
|
||||
});
|
||||
32
packages/ws/__tests__/util/SimpleIdentifyThrottler.test.ts
Normal file
32
packages/ws/__tests__/util/SimpleIdentifyThrottler.test.ts
Normal file
@@ -0,0 +1,32 @@
|
||||
import { setTimeout as sleep } from 'node:timers/promises';
|
||||
import { expect, test, vi, type Mock } from 'vitest';
|
||||
import { SimpleIdentifyThrottler } from '../../src/index.js';
|
||||
|
||||
vi.mock('node:timers/promises', () => ({
|
||||
setTimeout: vi.fn(),
|
||||
}));
|
||||
|
||||
const throttler = new SimpleIdentifyThrottler(2);
|
||||
|
||||
vi.useFakeTimers();
|
||||
|
||||
const NOW = vi.fn().mockReturnValue(Date.now());
|
||||
global.Date.now = NOW;
|
||||
|
||||
test('basic case', async () => {
|
||||
// Those shouldn't wait since they're in different keys
|
||||
|
||||
await throttler.waitForIdentify(0);
|
||||
expect(sleep).not.toHaveBeenCalled();
|
||||
|
||||
await throttler.waitForIdentify(1);
|
||||
expect(sleep).not.toHaveBeenCalled();
|
||||
|
||||
// Those should wait
|
||||
|
||||
await throttler.waitForIdentify(2);
|
||||
expect(sleep).toHaveBeenCalledTimes(1);
|
||||
|
||||
await throttler.waitForIdentify(3);
|
||||
expect(sleep).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
Reference in New Issue
Block a user