Nomchy/discord.js
1
0
Fork 0
mirror of https://github.com/discordjs/discord.js.git synced 2026-03-18 04:23:31 +01:00
Code Issues Packages Projects Releases Wiki Activity

refactor: abstract identify throttling and correct max_concurrency handling (#9375)

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files

30
packages/ws/__tests__/strategy/WorkerShardingStrategy.test.ts
View File

@@ -57,9 +57,9 @@ vi.mock('node:worker_threads', async () => {
this.emit('online');
// same deal here
setImmediate(() => {
const message = {
const message: WorkerReceivePayload = {
op: WorkerReceivePayloadOp.WorkerReady,
} satisfies WorkerReceivePayload;
};
this.emit('message', message);
});
});
@@ -68,39 +68,39 @@ vi.mock('node:worker_threads', async () => {
public postMessage(message: WorkerSendPayload) {
switch (message.op) {
case WorkerSendPayloadOp.Connect: {
const response = {
const response: WorkerReceivePayload = {
op: WorkerReceivePayloadOp.Connected,
shardId: message.shardId,
} satisfies WorkerReceivePayload;
};
this.emit('message', response);
break;
}
case WorkerSendPayloadOp.Destroy: {
const response = {
const response: WorkerReceivePayload = {
op: WorkerReceivePayloadOp.Destroyed,
shardId: message.shardId,
} satisfies WorkerReceivePayload;
};
this.emit('message', response);
break;
}
case WorkerSendPayloadOp.Send: {
if (message.payload.op === GatewayOpcodes.RequestGuildMembers) {
const response = {
const response: WorkerReceivePayload = {
op: WorkerReceivePayloadOp.Event,
shardId: message.shardId,
event: WebSocketShardEvents.Dispatch,
data: memberChunkData,
} satisfies WorkerReceivePayload;
};
this.emit('message', response);
// Fetch session info
const sessionFetch = {
const sessionFetch: WorkerReceivePayload = {
op: WorkerReceivePayloadOp.RetrieveSessionInfo,
shardId: message.shardId,
nonce: Math.random(),
} satisfies WorkerReceivePayload;
};
this.emit('message', sessionFetch);
}
@@ -111,16 +111,16 @@ vi.mock('node:worker_threads', async () => {
case WorkerSendPayloadOp.SessionInfoResponse: {
message.session ??= sessionInfo;
const session = {
const session: WorkerReceivePayload = {
op: WorkerReceivePayloadOp.UpdateSessionInfo,
shardId: message.session.shardId,
session: { ...message.session, sequence: message.session.sequence + 1 },
} satisfies WorkerReceivePayload;
};
this.emit('message', session);
break;
}
case WorkerSendPayloadOp.ShardCanIdentify: {
case WorkerSendPayloadOp.ShardIdentifyResponse: {
break;
}
@@ -198,10 +198,10 @@ test('spawn, connect, send a message, session info, and destroy', async () => {
expect.objectContaining({ workerData: expect.objectContaining({ shardIds: [0, 1] }) }),
);
const payload = {
const payload: GatewaySendPayload = {
op: GatewayOpcodes.RequestGuildMembers,
d: { guild_id: '123', limit: 0, query: '' },
} satisfies GatewaySendPayload;
};
await manager.send(0, payload);
expect(mockSend).toHaveBeenCalledWith(0, payload);
expect(managerEmitSpy).toHaveBeenCalledWith(WebSocketShardEvents.Dispatch, {

46
packages/ws/__tests__/util/IdentifyThrottler.test.ts
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();
});

32
packages/ws/__tests__/util/SimpleIdentifyThrottler.test.ts Normal file
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);
});