mirror of
https://github.com/discordjs/discord.js.git
synced 2026-03-15 19:13:31 +01:00
refactor(WebSocketShard): identify throttling (#8888)
* refactor(WebSocketShard): identify throttling * chore: add worker handling * refactor: worker handling * chore: update tests * chore: use satisfies where applicable * chore: add informative comment * chore: apply suggestions * refactor(SimpleContextFetchingStrategy): support multiple managers
This commit is contained in:
@@ -18,6 +18,7 @@ export interface IContextFetchingStrategy {
|
||||
readonly options: FetchingStrategyOptions;
|
||||
retrieveSessionInfo(shardId: number): Awaitable<SessionInfo | null>;
|
||||
updateSessionInfo(shardId: number, sessionInfo: SessionInfo | null): Awaitable<void>;
|
||||
waitForIdentify(): Promise<void>;
|
||||
}
|
||||
|
||||
export async function managerToFetchingStrategyOptions(manager: WebSocketManager): Promise<FetchingStrategyOptions> {
|
||||
|
||||
@@ -1,8 +1,28 @@
|
||||
import { IdentifyThrottler } from '../../utils/IdentifyThrottler.js';
|
||||
import type { SessionInfo, WebSocketManager } from '../../ws/WebSocketManager.js';
|
||||
import type { FetchingStrategyOptions, IContextFetchingStrategy } from './IContextFetchingStrategy.js';
|
||||
|
||||
export class SimpleContextFetchingStrategy implements IContextFetchingStrategy {
|
||||
public constructor(private readonly manager: WebSocketManager, public readonly options: FetchingStrategyOptions) {}
|
||||
// This strategy assumes every shard is running under the same process - therefore we need a single
|
||||
// IdentifyThrottler per manager.
|
||||
private static throttlerCache = new WeakMap<WebSocketManager, IdentifyThrottler>();
|
||||
|
||||
private static ensureThrottler(manager: WebSocketManager): IdentifyThrottler {
|
||||
const existing = SimpleContextFetchingStrategy.throttlerCache.get(manager);
|
||||
if (existing) {
|
||||
return existing;
|
||||
}
|
||||
|
||||
const throttler = new IdentifyThrottler(manager);
|
||||
SimpleContextFetchingStrategy.throttlerCache.set(manager, throttler);
|
||||
return throttler;
|
||||
}
|
||||
|
||||
private readonly throttler: IdentifyThrottler;
|
||||
|
||||
public constructor(private readonly manager: WebSocketManager, public readonly options: FetchingStrategyOptions) {
|
||||
this.throttler = SimpleContextFetchingStrategy.ensureThrottler(manager);
|
||||
}
|
||||
|
||||
public async retrieveSessionInfo(shardId: number): Promise<SessionInfo | null> {
|
||||
return this.manager.options.retrieveSessionInfo(shardId);
|
||||
@@ -11,4 +31,8 @@ export class SimpleContextFetchingStrategy implements IContextFetchingStrategy {
|
||||
public updateSessionInfo(shardId: number, sessionInfo: SessionInfo | null) {
|
||||
return this.manager.options.updateSessionInfo(shardId, sessionInfo);
|
||||
}
|
||||
|
||||
public async waitForIdentify(): Promise<void> {
|
||||
await this.throttler.waitForIdentify();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,6 +12,8 @@ import type { FetchingStrategyOptions, IContextFetchingStrategy } from './IConte
|
||||
export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
|
||||
private readonly sessionPromises = new Collection<number, (session: SessionInfo | null) => void>();
|
||||
|
||||
private readonly waitForIdentifyPromises = new Collection<number, () => void>();
|
||||
|
||||
public constructor(public readonly options: FetchingStrategyOptions) {
|
||||
if (isMainThread) {
|
||||
throw new Error('Cannot instantiate WorkerContextFetchingStrategy on the main thread');
|
||||
@@ -19,20 +21,24 @@ export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
|
||||
|
||||
parentPort!.on('message', (payload: WorkerSendPayload) => {
|
||||
if (payload.op === WorkerSendPayloadOp.SessionInfoResponse) {
|
||||
const resolve = this.sessionPromises.get(payload.nonce);
|
||||
resolve?.(payload.session);
|
||||
this.sessionPromises.get(payload.nonce)?.(payload.session);
|
||||
this.sessionPromises.delete(payload.nonce);
|
||||
}
|
||||
|
||||
if (payload.op === WorkerSendPayloadOp.ShardCanIdentify) {
|
||||
this.waitForIdentifyPromises.get(payload.nonce)?.();
|
||||
this.waitForIdentifyPromises.delete(payload.nonce);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
public async retrieveSessionInfo(shardId: number): Promise<SessionInfo | null> {
|
||||
const nonce = Math.random();
|
||||
const payload: WorkerRecievePayload = {
|
||||
const payload = {
|
||||
op: WorkerRecievePayloadOp.RetrieveSessionInfo,
|
||||
shardId,
|
||||
nonce,
|
||||
};
|
||||
} satisfies WorkerRecievePayload;
|
||||
// eslint-disable-next-line no-promise-executor-return
|
||||
const promise = new Promise<SessionInfo | null>((resolve) => this.sessionPromises.set(nonce, resolve));
|
||||
parentPort!.postMessage(payload);
|
||||
@@ -40,11 +46,23 @@ export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
|
||||
}
|
||||
|
||||
public updateSessionInfo(shardId: number, sessionInfo: SessionInfo | null) {
|
||||
const payload: WorkerRecievePayload = {
|
||||
const payload = {
|
||||
op: WorkerRecievePayloadOp.UpdateSessionInfo,
|
||||
shardId,
|
||||
session: sessionInfo,
|
||||
};
|
||||
} satisfies WorkerRecievePayload;
|
||||
parentPort!.postMessage(payload);
|
||||
}
|
||||
|
||||
public async waitForIdentify(): Promise<void> {
|
||||
const nonce = Math.random();
|
||||
const payload = {
|
||||
op: WorkerRecievePayloadOp.WaitForIdentify,
|
||||
nonce,
|
||||
} satisfies WorkerRecievePayload;
|
||||
// eslint-disable-next-line no-promise-executor-return
|
||||
const promise = new Promise<void>((resolve) => this.waitForIdentifyPromises.set(nonce, resolve));
|
||||
parentPort!.postMessage(payload);
|
||||
return promise;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user