From 442752319c7983d81e1767d4897c21ab68b0491c Mon Sep 17 00:00:00 2001 From: MARCROCK22 <57925328+MARCROCK22@users.noreply.github.com> Date: Wed, 20 Nov 2024 12:06:23 -0400 Subject: [PATCH] feat: custom worker events (#298) * chore: apply formatting * feat: custom worker events * fix: xd --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/client/workerclient.ts | 13 +++- src/commands/applications/shared.ts | 2 + src/common/shorters/soundboard.ts | 84 +++++++++++++------------- src/websocket/discord/shared.ts | 3 + src/websocket/discord/worker.ts | 20 +++++- src/websocket/discord/workermanager.ts | 41 ++++++++++--- 6 files changed, 109 insertions(+), 54 deletions(-) diff --git a/src/client/workerclient.ts b/src/client/workerclient.ts index 0968500..2a8bc1f 100644 --- a/src/client/workerclient.ts +++ b/src/client/workerclient.ts @@ -7,7 +7,7 @@ import type { GatewayDispatchPayload, GatewaySendPayload } from '../types'; import { Shard, type ShardManagerOptions, type WorkerData, properties } from '../websocket'; import type { WorkerDisconnectedAllShardsResharding, - WorkerMessage, + WorkerMessages, WorkerReady, WorkerReadyResharding, WorkerReceivePayload, @@ -78,6 +78,14 @@ export class WorkerClient extends BaseClient { if (options?.postMessage) { this.postMessage = options.postMessage; } + + if (this.options.handleManagerMessages) { + const oldFn = this.handleManagerMessages.bind(this); + this.handleManagerMessages = async message => { + await this.options.handleManagerMessages!(message); + return oldFn(message); + }; + } } get workerId() { @@ -157,7 +165,7 @@ export class WorkerClient extends BaseClient { } } - postMessage(body: WorkerMessage): unknown { + postMessage(body: WorkerMessages): unknown { if (manager) return manager.postMessage(body); return process.send!(body); } @@ -593,4 +601,5 @@ interface WorkerClientOptions extends BaseClientOptions { postMessage?: (body: unknown) => unknown; /** can have perfomance issues in big bots if the client sends every event, specially in startup (false by default) */ sendPayloadToParent?: boolean; + handleManagerMessages?(message: ManagerMessages): any; } diff --git a/src/commands/applications/shared.ts b/src/commands/applications/shared.ts index ad3f426..02048e0 100644 --- a/src/commands/applications/shared.ts +++ b/src/commands/applications/shared.ts @@ -27,6 +27,8 @@ export interface DefaultLocale {} export interface ExtendContext {} export interface ExtraProps {} export interface UsingClient extends BaseClient {} +export interface CustomWorkerClientEvents {} +export interface CustomWorkerManagerEvents {} export interface ExtendedRC {} export interface ExtendedRCLocations {} export type ParseClient = T; diff --git a/src/common/shorters/soundboard.ts b/src/common/shorters/soundboard.ts index fd6b22b..c656949 100644 --- a/src/common/shorters/soundboard.ts +++ b/src/common/shorters/soundboard.ts @@ -1,42 +1,42 @@ -import type { - RESTPatchAPIGuildSoundboardSound, - RESTPostAPIGuildSoundboardSound, - RESTPostAPISendSoundboardSound, -} from '../../types'; -import { BaseShorter } from './base'; - -export class SoundboardShorter extends BaseShorter { - getDefaults() { - return this.client.proxy['soundboard-default-sounds'].get(); - } - - send(channelId: string, body: RESTPostAPISendSoundboardSound) { - return this.client.proxy.channels(channelId)['send-soundboard-sound'].post({ - body, - }); - } - - list(guildId: string) { - return this.client.proxy.guilds(guildId)['soundboard-sounds'].get(); - } - - get(guildId: string, soundID: string) { - return this.client.proxy.guilds(guildId)['soundboard-sounds'](soundID).get(); - } - - create(guildId: string, body: RESTPostAPIGuildSoundboardSound) { - return this.client.proxy.guilds(guildId)['soundboard-sounds'].post({ - body, - }); - } - - edit(guildId: string, soundID: string, body: RESTPatchAPIGuildSoundboardSound) { - return this.client.proxy.guilds(guildId)['soundboard-sounds'](soundID).patch({ - body, - }); - } - - delete(guildId: string, soundID: string, reason?: string) { - return this.client.proxy.guilds(guildId)['soundboard-sounds'](soundID).delete({ reason }); - } -} +import type { + RESTPatchAPIGuildSoundboardSound, + RESTPostAPIGuildSoundboardSound, + RESTPostAPISendSoundboardSound, +} from '../../types'; +import { BaseShorter } from './base'; + +export class SoundboardShorter extends BaseShorter { + getDefaults() { + return this.client.proxy['soundboard-default-sounds'].get(); + } + + send(channelId: string, body: RESTPostAPISendSoundboardSound) { + return this.client.proxy.channels(channelId)['send-soundboard-sound'].post({ + body, + }); + } + + list(guildId: string) { + return this.client.proxy.guilds(guildId)['soundboard-sounds'].get(); + } + + get(guildId: string, soundID: string) { + return this.client.proxy.guilds(guildId)['soundboard-sounds'](soundID).get(); + } + + create(guildId: string, body: RESTPostAPIGuildSoundboardSound) { + return this.client.proxy.guilds(guildId)['soundboard-sounds'].post({ + body, + }); + } + + edit(guildId: string, soundID: string, body: RESTPatchAPIGuildSoundboardSound) { + return this.client.proxy.guilds(guildId)['soundboard-sounds'](soundID).patch({ + body, + }); + } + + delete(guildId: string, soundID: string, reason?: string) { + return this.client.proxy.guilds(guildId)['soundboard-sounds'](soundID).delete({ reason }); + } +} diff --git a/src/websocket/discord/shared.ts b/src/websocket/discord/shared.ts index 1324dbd..c90ae4c 100644 --- a/src/websocket/discord/shared.ts +++ b/src/websocket/discord/shared.ts @@ -6,6 +6,7 @@ import type { GatewayPresenceUpdateData, } from '../../types'; import type { IdentifyProperties } from '../constants'; +import type { WorkerMessages } from './worker'; export interface ShardManagerOptions extends ShardDetails { /** Important data which is used by the manager to connect shards to the gateway. */ @@ -69,6 +70,8 @@ export interface WorkerManagerOptions extends Omit>; } diff --git a/src/websocket/discord/worker.ts b/src/websocket/discord/worker.ts index 08a3ef7..6bd447d 100644 --- a/src/websocket/discord/worker.ts +++ b/src/websocket/discord/worker.ts @@ -1,4 +1,5 @@ -import type { ApiRequestOptions, HttpMethods } from '../..'; +import type { ApiRequestOptions, CustomWorkerClientEvents, HttpMethods } from '../..'; +import type { Identify } from '../../common'; import type { GatewayDispatchPayload } from '../../types'; export interface WorkerShardInfo { @@ -86,7 +87,7 @@ export type WorkerSendToWorkerEval = CreateWorkerMessage< } >; -export type WorkerMessage = +export type BaseWorkerMessage = | WorkerRequestConnect | WorkerReceivePayload | WorkerSendResultPayload @@ -103,3 +104,18 @@ export type WorkerMessage = | WorkerRequestConnectResharding | WorkerReadyResharding | WorkerDisconnectedAllShardsResharding; + +export type CustomWorkerClientMessages = { + [K in keyof CustomWorkerClientEvents]: Identify< + { + type: K; + workerId: number; + } & Identify + >; +}; + +export type WorkerMessages = + | { + [K in BaseWorkerMessage['type']]: Identify>; + }[BaseWorkerMessage['type']] + | CustomWorkerClientMessages[keyof CustomWorkerClientMessages]; diff --git a/src/websocket/discord/workermanager.ts b/src/websocket/discord/workermanager.ts index fd249cb..478c27f 100644 --- a/src/websocket/discord/workermanager.ts +++ b/src/websocket/discord/workermanager.ts @@ -1,16 +1,16 @@ import cluster, { type Worker as ClusterWorker } from 'node:cluster'; import { type UUID, randomUUID } from 'node:crypto'; import type { Worker as WorkerThreadsWorker } from 'node:worker_threads'; -import { ApiHandler, Logger, type UsingClient, type WorkerClient } from '../..'; +import { ApiHandler, type CustomWorkerManagerEvents, Logger, type UsingClient, type WorkerClient } from '../..'; import { type Adapter, MemoryAdapter } from '../../cache'; import { BaseClient, type InternalRuntimeConfig } from '../../client/base'; -import { BASE_HOST, type MakePartial, MergeOptions, lazyLoadPackage } from '../../common'; +import { BASE_HOST, type Identify, type MakePartial, MergeOptions, lazyLoadPackage } from '../../common'; import type { GatewayPresenceUpdateData, GatewaySendPayload, RESTGetAPIGatewayBotResult } from '../../types'; import { WorkerManagerDefaults, properties } from '../constants'; import { DynamicBucket } from '../structures'; import { ConnectQueue } from '../structures/timeout'; import type { ShardOptions, WorkerData, WorkerManagerOptions } from './shared'; -import type { WorkerInfo, WorkerMessage, WorkerShardInfo } from './worker'; +import type { WorkerInfo, WorkerMessages, WorkerShardInfo } from './worker'; export class WorkerManager extends Map< number, @@ -46,7 +46,7 @@ export class WorkerManager extends Map< return chunks; } - options: MakePartial, 'adapter'>; + options: MakePartial, 'adapter' | 'handleWorkerMessage' | 'handlePayload'>; debugger?: Logger; connectQueue!: ConnectQueue; workerQueue: (() => void)[] = []; @@ -57,13 +57,24 @@ export class WorkerManager extends Map< private _info?: RESTGetAPIGatewayBotResult; constructor( - options: Omit, 'resharding'> & { + options: Omit< + MakePartial, + 'resharding' + > & { resharding?: MakePartial, 'getInfo'>; }, ) { super(); this.options = options as WorkerManager['options']; this.cacheAdapter = new MemoryAdapter(); + + if (this.options.handleWorkerMessage) { + const oldFn = this.handleWorkerMessage.bind(this); + this.handleWorkerMessage = async message => { + await this.options.handleWorkerMessage!(message); + return oldFn(message); + }; + } } setCache(adapter: Adapter) { @@ -241,7 +252,7 @@ export class WorkerManager extends Map< }); } - async handleWorkerMessage(message: WorkerMessage) { + async handleWorkerMessage(message: WorkerMessages) { switch (message.type) { case 'WORKER_READY_RESHARDING': { @@ -340,7 +351,7 @@ export class WorkerManager extends Map< } break; case 'RECEIVE_PAYLOAD': - await this.options.handlePayload(message.shardId, message.workerId, message.payload); + await this.options.handlePayload?.(message.shardId, message.workerId, message.payload); break; case 'RESULT_PAYLOAD': { @@ -687,7 +698,7 @@ export type ManagerSendEvalResponse = CreateManagerMessage< } >; -export type ManagerMessages = +export type BaseManagerMessages = | ManagerAllowConnect | ManagerSpawnShards | ManagerSendPayload @@ -704,3 +715,17 @@ export type ManagerMessages = | DisconnectAllShardsResharding | ConnnectAllShardsResharding | ManagerExecuteEval; + +export type CustomManagerMessages = { + [K in keyof CustomWorkerManagerEvents]: Identify< + { + type: K; + } & Identify + >; +}; + +export type ManagerMessages = + | { + [K in BaseManagerMessages['type']]: Identify>; + }[BaseManagerMessages['type']] + | CustomManagerMessages[keyof CustomManagerMessages];