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:
DD
2023-04-14 23:26:37 +03:00
committed by GitHub
parent cac3c07729
commit 02dfaf1aa2
16 changed files with 279 additions and 161 deletions

View File

@@ -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();
});

View 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);
});