Nomchy/discord.js
1
0
Fork 0
mirror of https://github.com/discordjs/discord.js.git synced 2026-03-19 21:13:30 +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

1
packages/ws/src/strategies/sharding/SimpleShardingStrategy.ts
View File

@@ -23,6 +23,7 @@ export class SimpleShardingStrategy implements IShardingStrategy {
*/
public async spawn(shardIds: number[]) {
const strategyOptions = await managerToFetchingStrategyOptions(this.manager);
for (const shardId of shardIds) {
const strategy = new SimpleContextFetchingStrategy(this.manager, strategyOptions);
const shard = new WebSocketShard(strategy, shardId);

65
packages/ws/src/strategies/sharding/WorkerShardingStrategy.ts
View File

@@ -3,7 +3,7 @@ import { join, isAbsolute, resolve } from 'node:path';
import { Worker } from 'node:worker_threads';
import { Collection } from '@discordjs/collection';
import type { GatewaySendPayload } from 'discord-api-types/v10';
import { IdentifyThrottler } from '../../utils/IdentifyThrottler.js';
import type { IIdentifyThrottler } from '../../throttling/IIdentifyThrottler';
import type { SessionInfo, WebSocketManager } from '../../ws/WebSocketManager';
import type { WebSocketShardDestroyOptions, WebSocketShardEvents, WebSocketShardStatus } from '../../ws/WebSocketShard';
import { managerToFetchingStrategyOptions, type FetchingStrategyOptions } from '../context/IContextFetchingStrategy.js';
@@ -18,14 +18,14 @@ export enum WorkerSendPayloadOp {
Destroy,
Send,
SessionInfoResponse,
ShardCanIdentify,
ShardIdentifyResponse,
FetchStatus,
}
export type WorkerSendPayload =
| { nonce: number; ok: boolean; op: WorkerSendPayloadOp.ShardIdentifyResponse }
| { nonce: number; op: WorkerSendPayloadOp.FetchStatus; shardId: number }
| { nonce: number; op: WorkerSendPayloadOp.SessionInfoResponse; session: SessionInfo | null }
| { nonce: number; op: WorkerSendPayloadOp.ShardCanIdentify }
| { op: WorkerSendPayloadOp.Connect; shardId: number }
| { op: WorkerSendPayloadOp.Destroy; options?: WebSocketShardDestroyOptions; shardId: number }
| { op: WorkerSendPayloadOp.Send; payload: GatewaySendPayload; shardId: number };
@@ -39,14 +39,16 @@ export enum WorkerReceivePayloadOp {
WaitForIdentify,
FetchStatusResponse,
WorkerReady,
CancelIdentify,
}
export type WorkerReceivePayload =
// Can't seem to get a type-safe union based off of the event, so I'm sadly leaving data as any for now
| { data: any; event: WebSocketShardEvents; op: WorkerReceivePayloadOp.Event; shardId: number }
| { nonce: number; op: WorkerReceivePayloadOp.CancelIdentify }
| { nonce: number; op: WorkerReceivePayloadOp.FetchStatusResponse; status: WebSocketShardStatus }
| { nonce: number; op: WorkerReceivePayloadOp.RetrieveSessionInfo; shardId: number }
| { nonce: number; op: WorkerReceivePayloadOp.WaitForIdentify }
| { nonce: number; op: WorkerReceivePayloadOp.WaitForIdentify; shardId: number }
| { op: WorkerReceivePayloadOp.Connected; shardId: number }
| { op: WorkerReceivePayloadOp.Destroyed; shardId: number }
| { op: WorkerReceivePayloadOp.UpdateSessionInfo; session: SessionInfo | null; shardId: number }
@@ -84,11 +86,12 @@ export class WorkerShardingStrategy implements IShardingStrategy {
private readonly fetchStatusPromises = new Collection<number, (status: WebSocketShardStatus) => void>();
private readonly throttler: IdentifyThrottler;
private readonly waitForIdentifyControllers = new Collection<number, AbortController>();
private throttler?: IIdentifyThrottler;
public constructor(manager: WebSocketManager, options: WorkerShardingStrategyOptions) {
this.manager = manager;
this.throttler = new IdentifyThrottler(manager);
this.options = options;
}
@@ -122,10 +125,10 @@ export class WorkerShardingStrategy implements IShardingStrategy {
const promises = [];
for (const [shardId, worker] of this.#workerByShardId.entries()) {
const payload = {
const payload: WorkerSendPayload = {
op: WorkerSendPayloadOp.Connect,
shardId,
} satisfies WorkerSendPayload;
};
// eslint-disable-next-line no-promise-executor-return
const promise = new Promise<void>((resolve) => this.connectPromises.set(shardId, resolve));
@@ -143,11 +146,11 @@ export class WorkerShardingStrategy implements IShardingStrategy {
const promises = [];
for (const [shardId, worker] of this.#workerByShardId.entries()) {
const payload = {
const payload: WorkerSendPayload = {
op: WorkerSendPayloadOp.Destroy,
shardId,
options,
} satisfies WorkerSendPayload;
};
promises.push(
// eslint-disable-next-line no-promise-executor-return, promise/prefer-await-to-then
@@ -171,11 +174,11 @@ export class WorkerShardingStrategy implements IShardingStrategy {
throw new Error(`No worker found for shard ${shardId}`);
}
const payload = {
const payload: WorkerSendPayload = {
op: WorkerSendPayloadOp.Send,
shardId,
payload: data,
} satisfies WorkerSendPayload;
};
worker.postMessage(payload);
}
@@ -187,11 +190,11 @@ export class WorkerShardingStrategy implements IShardingStrategy {
for (const [shardId, worker] of this.#workerByShardId.entries()) {
const nonce = Math.random();
const payload = {
const payload: WorkerSendPayload = {
op: WorkerSendPayloadOp.FetchStatus,
shardId,
nonce,
} satisfies WorkerSendPayload;
};
// eslint-disable-next-line no-promise-executor-return
const promise = new Promise<WebSocketShardStatus>((resolve) => this.fetchStatusPromises.set(nonce, resolve));
@@ -297,10 +300,21 @@ export class WorkerShardingStrategy implements IShardingStrategy {
}
case WorkerReceivePayloadOp.WaitForIdentify: {
await this.throttler.waitForIdentify();
const throttler = await this.ensureThrottler();
// If this rejects it means we aborted, in which case we reply elsewhere.
try {
const controller = new AbortController();
this.waitForIdentifyControllers.set(payload.nonce, controller);
await throttler.waitForIdentify(payload.shardId, controller.signal);
} catch {
return;
}
const response: WorkerSendPayload = {
op: WorkerSendPayloadOp.ShardCanIdentify,
op: WorkerSendPayloadOp.ShardIdentifyResponse,
nonce: payload.nonce,
ok: true,
};
worker.postMessage(response);
break;
@@ -315,6 +329,25 @@ export class WorkerShardingStrategy implements IShardingStrategy {
case WorkerReceivePayloadOp.WorkerReady: {
break;
}
case WorkerReceivePayloadOp.CancelIdentify: {
this.waitForIdentifyControllers.get(payload.nonce)?.abort();
this.waitForIdentifyControllers.delete(payload.nonce);
const response: WorkerSendPayload = {
op: WorkerSendPayloadOp.ShardIdentifyResponse,
nonce: payload.nonce,
ok: false,
};
worker.postMessage(response);
break;
}
}
}
private async ensureThrottler(): Promise<IIdentifyThrottler> {
this.throttler ??= await this.manager.options.buildIdentifyThrottler(this.manager);
return this.throttler;
}
}