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:
DD
2022-12-02 15:04:09 +02:00
committed by GitHub
parent 3fca638a84
commit 8f552a0e17
9 changed files with 113 additions and 38 deletions

View File

@@ -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> {

View File

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

View File

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