refactor(core): allow easily supporting redis gateway

didinele · didinele · commit 73ba97a3afbd · 2023-05-08T12:17:40.000+03:00
BREAKING CHANGE: the client now requires a shardCount prop
diff --git a/packages/core/README.md b/packages/core/README.md
@@ -48,7 +48,7 @@ const gateway = new WebSocketManager({
 });
 
 // Create a client to emit relevant events.
-const client = new Client({ rest, gateway });
+const client = new Client({ rest, gateway, shardCount: await gateway.getShardCount() });
 
 // Listen for interactions
 // Each event contains an `api` prop along with the event data that allows you to interface with the Discord REST API
diff --git a/packages/core/package.json b/packages/core/package.json
@@ -62,6 +62,7 @@
 		"discord-api-types": "^0.37.41"
 	},
 	"devDependencies": {
+		"@discordjs/brokers": "workspace:^",
 		"@favware/cliff-jumper": "^2.0.0",
 		"@microsoft/api-extractor": "^7.34.8",
 		"@types/node": "18.16.5",
diff --git a/packages/core/src/client.ts b/packages/core/src/client.ts
@@ -4,75 +4,16 @@ import { calculateShardId } from '@discordjs/util';
 import { WebSocketShardEvents } from '@discordjs/ws';
 import { DiscordSnowflake } from '@sapphire/snowflake';
 import { AsyncEventEmitter } from '@vladfrangu/async_event_emitter';
-import {
-	GatewayDispatchEvents,
-	GatewayOpcodes,
-	type APIGuildMember,
-	type GatewayAutoModerationActionExecutionDispatchData,
-	type GatewayAutoModerationRuleCreateDispatchData,
-	type GatewayAutoModerationRuleDeleteDispatchData,
-	type GatewayAutoModerationRuleUpdateDispatchData,
-	type GatewayChannelCreateDispatchData,
-	type GatewayChannelDeleteDispatchData,
-	type GatewayChannelPinsUpdateDispatchData,
-	type GatewayChannelUpdateDispatchData,
-	type GatewayGuildAuditLogEntryCreateDispatchData,
-	type GatewayGuildBanAddDispatchData,
-	type GatewayGuildBanRemoveDispatchData,
-	type GatewayGuildCreateDispatchData,
-	type GatewayGuildDeleteDispatchData,
-	type GatewayGuildEmojisUpdateDispatchData,
-	type GatewayGuildIntegrationsUpdateDispatchData,
-	type GatewayGuildMemberAddDispatchData,
-	type GatewayGuildMemberRemoveDispatchData,
-	type GatewayGuildMemberUpdateDispatchData,
-	type GatewayGuildMembersChunkDispatchData,
-	type GatewayGuildRoleCreateDispatchData,
-	type GatewayGuildRoleDeleteDispatchData,
-	type GatewayGuildRoleUpdateDispatchData,
-	type GatewayGuildScheduledEventCreateDispatchData,
-	type GatewayGuildScheduledEventDeleteDispatchData,
-	type GatewayGuildScheduledEventUpdateDispatchData,
-	type GatewayGuildScheduledEventUserAddDispatchData,
-	type GatewayGuildScheduledEventUserRemoveDispatchData,
-	type GatewayGuildStickersUpdateDispatchData,
-	type GatewayGuildUpdateDispatchData,
-	type GatewayIntegrationCreateDispatchData,
-	type GatewayIntegrationDeleteDispatchData,
-	type GatewayIntegrationUpdateDispatchData,
-	type GatewayInteractionCreateDispatchData,
-	type GatewayInviteCreateDispatchData,
-	type GatewayInviteDeleteDispatchData,
-	type GatewayMessageCreateDispatchData,
-	type GatewayMessageDeleteBulkDispatchData,
-	type GatewayMessageDeleteDispatchData,
-	type GatewayMessageReactionAddDispatchData,
-	type GatewayMessageReactionRemoveAllDispatchData,
-	type GatewayMessageReactionRemoveDispatchData,
-	type GatewayMessageReactionRemoveEmojiDispatchData,
-	type GatewayMessageUpdateDispatchData,
-	type GatewayPresenceUpdateData,
-	type GatewayPresenceUpdateDispatchData,
-	type GatewayReadyDispatchData,
-	type GatewayRequestGuildMembersData,
-	type GatewayStageInstanceCreateDispatchData,
-	type GatewayStageInstanceDeleteDispatchData,
-	type GatewayStageInstanceUpdateDispatchData,
-	type GatewayThreadCreateDispatchData,
-	type GatewayThreadDeleteDispatchData,
-	type GatewayThreadListSyncDispatchData,
-	type GatewayThreadMemberUpdateDispatchData,
-	type GatewayThreadMembersUpdateDispatchData,
-	type GatewayThreadUpdateDispatchData,
-	type GatewayTypingStartDispatchData,
-	type GatewayUserUpdateDispatchData,
-	type GatewayVoiceServerUpdateDispatchData,
-	type GatewayVoiceStateUpdateData,
-	type GatewayVoiceStateUpdateDispatchData,
-	type GatewayWebhooksUpdateDispatchData,
+import { GatewayDispatchEvents, GatewayOpcodes } from 'discord-api-types/v10';
+import type {
+	GatewayDispatchPayload,
+	APIGuildMember,
+	GatewayRequestGuildMembersData,
+	GatewayPresenceUpdateData,
+	GatewayVoiceStateUpdateData,
 } from 'discord-api-types/v10';
-import type { Gateway } from './Gateway.js';
 import { API } from './api/index.js';
+import type { Gateway } from './gateway/Gateway.js';
 
 export interface IntrinsicProps {
 	/**
@@ -89,98 +30,43 @@ export interface WithIntrinsicProps<T> extends IntrinsicProps {
 	data: T;
 }
 
-export interface MappedEvents {
-	[GatewayDispatchEvents.AutoModerationActionExecution]: [
-		WithIntrinsicProps<GatewayAutoModerationActionExecutionDispatchData>,
-	];
-	[GatewayDispatchEvents.AutoModerationRuleCreate]: [WithIntrinsicProps<GatewayAutoModerationRuleCreateDispatchData>];
-	[GatewayDispatchEvents.AutoModerationRuleDelete]: [WithIntrinsicProps<GatewayAutoModerationRuleDeleteDispatchData>];
-	[GatewayDispatchEvents.AutoModerationRuleUpdate]: [WithIntrinsicProps<GatewayAutoModerationRuleUpdateDispatchData>];
-	[GatewayDispatchEvents.ChannelCreate]: [WithIntrinsicProps<GatewayChannelCreateDispatchData>];
-	[GatewayDispatchEvents.ChannelDelete]: [WithIntrinsicProps<GatewayChannelDeleteDispatchData>];
-	[GatewayDispatchEvents.ChannelPinsUpdate]: [WithIntrinsicProps<GatewayChannelPinsUpdateDispatchData>];
-	[GatewayDispatchEvents.ChannelUpdate]: [WithIntrinsicProps<GatewayChannelUpdateDispatchData>];
-	[GatewayDispatchEvents.GuildAuditLogEntryCreate]: [WithIntrinsicProps<GatewayGuildAuditLogEntryCreateDispatchData>];
-	[GatewayDispatchEvents.GuildBanAdd]: [WithIntrinsicProps<GatewayGuildBanAddDispatchData>];
-	[GatewayDispatchEvents.GuildBanRemove]: [WithIntrinsicProps<GatewayGuildBanRemoveDispatchData>];
-	[GatewayDispatchEvents.GuildCreate]: [WithIntrinsicProps<GatewayGuildCreateDispatchData>];
-	[GatewayDispatchEvents.GuildDelete]: [WithIntrinsicProps<GatewayGuildDeleteDispatchData>];
-	[GatewayDispatchEvents.GuildEmojisUpdate]: [WithIntrinsicProps<GatewayGuildEmojisUpdateDispatchData>];
-	[GatewayDispatchEvents.GuildIntegrationsUpdate]: [WithIntrinsicProps<GatewayGuildIntegrationsUpdateDispatchData>];
-	[GatewayDispatchEvents.GuildMemberAdd]: [WithIntrinsicProps<GatewayGuildMemberAddDispatchData>];
-	[GatewayDispatchEvents.GuildMemberRemove]: [WithIntrinsicProps<GatewayGuildMemberRemoveDispatchData>];
-	[GatewayDispatchEvents.GuildMemberUpdate]: [WithIntrinsicProps<GatewayGuildMemberUpdateDispatchData>];
-	[GatewayDispatchEvents.GuildMembersChunk]: [WithIntrinsicProps<GatewayGuildMembersChunkDispatchData>];
-	[GatewayDispatchEvents.GuildRoleCreate]: [WithIntrinsicProps<GatewayGuildRoleCreateDispatchData>];
-	[GatewayDispatchEvents.GuildRoleDelete]: [WithIntrinsicProps<GatewayGuildRoleDeleteDispatchData>];
-	[GatewayDispatchEvents.GuildRoleUpdate]: [WithIntrinsicProps<GatewayGuildRoleUpdateDispatchData>];
-	[GatewayDispatchEvents.GuildScheduledEventCreate]: [WithIntrinsicProps<GatewayGuildScheduledEventCreateDispatchData>];
-	[GatewayDispatchEvents.GuildScheduledEventDelete]: [WithIntrinsicProps<GatewayGuildScheduledEventDeleteDispatchData>];
-	[GatewayDispatchEvents.GuildScheduledEventUpdate]: [WithIntrinsicProps<GatewayGuildScheduledEventUpdateDispatchData>];
-	[GatewayDispatchEvents.GuildScheduledEventUserAdd]: [
-		WithIntrinsicProps<GatewayGuildScheduledEventUserAddDispatchData>,
-	];
-	[GatewayDispatchEvents.GuildScheduledEventUserRemove]: [
-		WithIntrinsicProps<GatewayGuildScheduledEventUserRemoveDispatchData>,
-	];
-	[GatewayDispatchEvents.GuildStickersUpdate]: [WithIntrinsicProps<GatewayGuildStickersUpdateDispatchData>];
-	[GatewayDispatchEvents.GuildUpdate]: [WithIntrinsicProps<GatewayGuildUpdateDispatchData>];
-	[GatewayDispatchEvents.IntegrationCreate]: [WithIntrinsicProps<GatewayIntegrationCreateDispatchData>];
-	[GatewayDispatchEvents.IntegrationDelete]: [WithIntrinsicProps<GatewayIntegrationDeleteDispatchData>];
-	[GatewayDispatchEvents.IntegrationUpdate]: [WithIntrinsicProps<GatewayIntegrationUpdateDispatchData>];
-	[GatewayDispatchEvents.InteractionCreate]: [WithIntrinsicProps<GatewayInteractionCreateDispatchData>];
-	[GatewayDispatchEvents.InviteCreate]: [WithIntrinsicProps<GatewayInviteCreateDispatchData>];
-	[GatewayDispatchEvents.InviteDelete]: [WithIntrinsicProps<GatewayInviteDeleteDispatchData>];
-	[GatewayDispatchEvents.MessageCreate]: [WithIntrinsicProps<GatewayMessageCreateDispatchData>];
-	[GatewayDispatchEvents.MessageDelete]: [WithIntrinsicProps<GatewayMessageDeleteDispatchData>];
-	[GatewayDispatchEvents.MessageDeleteBulk]: [WithIntrinsicProps<GatewayMessageDeleteBulkDispatchData>];
-	[GatewayDispatchEvents.MessageReactionAdd]: [WithIntrinsicProps<GatewayMessageReactionAddDispatchData>];
-	[GatewayDispatchEvents.MessageReactionRemove]: [WithIntrinsicProps<GatewayMessageReactionRemoveDispatchData>];
-	[GatewayDispatchEvents.MessageReactionRemoveAll]: [WithIntrinsicProps<GatewayMessageReactionRemoveAllDispatchData>];
-	[GatewayDispatchEvents.MessageReactionRemoveEmoji]: [
-		WithIntrinsicProps<GatewayMessageReactionRemoveEmojiDispatchData>,
-	];
-	[GatewayDispatchEvents.MessageUpdate]: [WithIntrinsicProps<GatewayMessageUpdateDispatchData>];
-	[GatewayDispatchEvents.PresenceUpdate]: [WithIntrinsicProps<GatewayPresenceUpdateDispatchData>];
-	[GatewayDispatchEvents.Ready]: [WithIntrinsicProps<GatewayReadyDispatchData>];
-	[GatewayDispatchEvents.Resumed]: [WithIntrinsicProps<never>];
-	[GatewayDispatchEvents.StageInstanceCreate]: [WithIntrinsicProps<GatewayStageInstanceCreateDispatchData>];
-	[GatewayDispatchEvents.StageInstanceDelete]: [WithIntrinsicProps<GatewayStageInstanceDeleteDispatchData>];
-	[GatewayDispatchEvents.StageInstanceUpdate]: [WithIntrinsicProps<GatewayStageInstanceUpdateDispatchData>];
-	[GatewayDispatchEvents.ThreadCreate]: [WithIntrinsicProps<GatewayThreadCreateDispatchData>];
-	[GatewayDispatchEvents.ThreadDelete]: [WithIntrinsicProps<GatewayThreadDeleteDispatchData>];
-	[GatewayDispatchEvents.ThreadListSync]: [WithIntrinsicProps<GatewayThreadListSyncDispatchData>];
-	[GatewayDispatchEvents.ThreadMemberUpdate]: [WithIntrinsicProps<GatewayThreadMemberUpdateDispatchData>];
-	[GatewayDispatchEvents.ThreadMembersUpdate]: [WithIntrinsicProps<GatewayThreadMembersUpdateDispatchData>];
-	[GatewayDispatchEvents.ThreadUpdate]: [WithIntrinsicProps<GatewayThreadUpdateDispatchData>];
-	[GatewayDispatchEvents.TypingStart]: [WithIntrinsicProps<GatewayTypingStartDispatchData>];
-	[GatewayDispatchEvents.UserUpdate]: [WithIntrinsicProps<GatewayUserUpdateDispatchData>];
-	[GatewayDispatchEvents.VoiceServerUpdate]: [WithIntrinsicProps<GatewayVoiceServerUpdateDispatchData>];
-	[GatewayDispatchEvents.VoiceStateUpdate]: [WithIntrinsicProps<GatewayVoiceStateUpdateDispatchData>];
-	[GatewayDispatchEvents.WebhooksUpdate]: [WithIntrinsicProps<GatewayWebhooksUpdateDispatchData>];
-}
+// need this to be its own type for some reason, the compiler doesn't behave the same way if we in-line it
+type _DiscordEvents = {
+	[K in GatewayDispatchEvents]: GatewayDispatchPayload & {
+		t: K;
+	};
+};
 
-export type ManagerShardEventsMap = {
-	[K in keyof MappedEvents]: MappedEvents[K];
+export type DiscordEvents = {
+	// @ts-expect-error - unclear why this is an error, this behaves as expected
+	[K in keyof _DiscordEvents]: _DiscordEvents[K]['d'];
+};
+
+export type MappedEvents = {
+	[K in keyof DiscordEvents]: [WithIntrinsicProps<DiscordEvents[K]>];
 };
 
 export interface ClientOptions {
 	gateway: Gateway;
 	rest: REST;
+	shardCount: number;
 }
 
-export class Client extends AsyncEventEmitter<ManagerShardEventsMap> {
+export class Client extends AsyncEventEmitter<MappedEvents> {
 	public readonly rest: REST;
 
 	public readonly gateway: Gateway;
 
 	public readonly api: API;
 
-	public constructor({ rest, gateway }: ClientOptions) {
+	public readonly shardCount: number;
+
+	public constructor({ rest, gateway, shardCount }: ClientOptions) {
 		super();
 		this.rest = rest;
 		this.gateway = gateway;
 		this.api = new API(rest);
+		this.shardCount = shardCount;
 
 		this.gateway.on(WebSocketShardEvents.Dispatch, ({ data: dispatch, shardId }) => {
 			// @ts-expect-error event props can't be resolved properly, but they are correct
@@ -196,7 +82,7 @@ export class Client extends AsyncEventEmitter<ManagerShardEventsMap> {
 	 * @param timeout - The timeout for waiting for each guild members chunk event
 	 */
 	public async requestGuildMembers(options: GatewayRequestGuildMembersData, timeout = 10_000) {
-		const shardId = calculateShardId(options.guild_id, await this.gateway.getShardCount());
+		const shardId = calculateShardId(options.guild_id, this.shardCount);
 		const nonce = options.nonce ?? DiscordSnowflake.generate().toString();
 
 		const promise = new Promise<APIGuildMember[]>((resolve, reject) => {
@@ -241,7 +127,7 @@ export class Client extends AsyncEventEmitter<ManagerShardEventsMap> {
 	 * @param options - The options for updating the voice state
 	 */
 	public async updateVoiceState(options: GatewayVoiceStateUpdateData) {
-		const shardId = calculateShardId(options.guild_id, await this.gateway.getShardCount());
+		const shardId = calculateShardId(options.guild_id, this.shardCount);
 
 		await this.gateway.send(shardId, {
 			op: GatewayOpcodes.VoiceStateUpdate,
diff --git a/packages/core/src/gateway/Gateway.ts b/packages/core/src/gateway/Gateway.ts
@@ -6,10 +6,6 @@ import type { GatewaySendPayload } from 'discord-api-types/v10';
  * A Discord gateway-like interface that can be used to send & recieve events.
  */
 export interface Gateway {
-	/**
-	 * Gets how many shards your bot is running.
-	 */
-	getShardCount(): Awaitable<number>;
 	on(
 		event: WebSocketShardEvents.Dispatch,
 		listener: (...params: ManagerShardEventsMap[WebSocketShardEvents.Dispatch]) => Awaitable<void>,
diff --git a/packages/core/src/gateway/RedisGateway.ts b/packages/core/src/gateway/RedisGateway.ts
@@ -0,0 +1,43 @@
+import type { PubSubRedisBroker } from '@discordjs/brokers';
+import type { ManagerShardEventsMap, WebSocketShardEvents } from '@discordjs/ws';
+import { AsyncEventEmitter } from '@vladfrangu/async_event_emitter';
+import type { GatewaySendPayload, GatewayDispatchEvents } from 'discord-api-types/v10';
+import type { DiscordEvents } from '../client.js';
+import type { Gateway } from './Gateway.js';
+
+interface BrokerIntrinsicProps {
+	shardId: number;
+}
+
+interface Events extends DiscordEvents {
+	gateway_send: GatewaySendPayload;
+}
+
+export type RedisBrokerDiscordEvents = {
+	[K in keyof Events]: BrokerIntrinsicProps & { payload: Events[K] };
+};
+
+export class RedisGateway
+	extends AsyncEventEmitter<{ dispatch: ManagerShardEventsMap[WebSocketShardEvents.Dispatch] }>
+	implements Gateway
+{
+	public constructor(private readonly broker: PubSubRedisBroker<RedisBrokerDiscordEvents>) {
+		super();
+	}
+
+	public async send(shardId: number, payload: GatewaySendPayload): Promise<void> {
+		await this.broker.publish('gateway_send', { payload, shardId });
+	}
+
+	public async init(group: string, events: GatewayDispatchEvents[]) {
+		for (const event of events) {
+			this.broker.on(event, ({ data: { payload, shardId }, ack }) => {
+				// @ts-expect-error - Union shenanigans
+				this.emit('dispatch', { shardId, data: payload });
+				void ack();
+			});
+		}
+
+		await this.broker.subscribe(group, events);
+	}
+}
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
@@ -1,6 +1,8 @@
 export * from './api/index.js';
-export * from './client.js';
+export * from './gateway/Gateway.js';
+export * from './gateway/RedisGateway.js';
 export * from './util/index.js';
+export * from './client.js';
 
 export * from 'discord-api-types/v10';
 
diff --git a/packages/redis-gateway/README.md b/packages/redis-gateway/README.md
@@ -70,14 +70,54 @@ await broker.publish('gateway_send', {
 		status: PresenceUpdateStatus.Online,
 	},
 });
+
+// if you were to start this process multiple times (e.g. multiple apps using 'work_balancing_group'),
+// they would automatically work balance those interaction create events
+await broker.subscribe('work_balancing_group', [GatewayDispatchEvents.InteractionCreate]);
 ```
 
-For TypeScript usage, you can pass in a gereric type to the `PubSubRedisBroker` to map out all the events,
-refer to [this container's implementation](https://github.com/discordjs/discord.js/tree/main/packages/redis-gateway/src/index.ts#L15) for reference.
+For TypeScript usage, you can pass in a gereric type to the `PubSubRedisBroker` to map out all the events, a mapped
+interface is available in `@discordjs/core` as `RedisBrokerDiscordEvents`.
+
+If you wish, you can also just use `@discordjs/core`:
+
+```ts
+import { REST } from '@discordjs/rest';
+import Redis from 'ioredis';
+import { PubSubRedisBroker } from '@discordjs/brokers';
+import {
+	GatewayDispatchEvents,
+	GatewayIntentBits,
+	InteractionType,
+	MessageFlags,
+	Client,
+	RedisGateway,
+} from '@discordjs/core';
+
+const rest = new REST({ version: '10' }).setToken(token);
+
+const redis = new Redis();
+const broker = new PubSubRedisBroker({ redisClient: redis });
+const gateway = new RedisGateway(broker);
+
+const client = new Client({
+	rest,
+	gateway,
+	// you can get this however you want, it's used for some calculations and should be your bot's TOTAL shard count
+	// across "clusters" or anything else.
+	shardCount: Number(process.env.SHARD_COUNT!),
+});
+
+client.on(GatewayDispatchEvents.InteractionCreate, async ({ data: interaction, api }) => {
+	if (interaction.type !== InteractionType.ApplicationCommand || interaction.data.name !== 'ping') {
+		return;
+	}
+
+	await api.interactions.reply(interaction.id, interaction.token, { content: 'Pong!', flags: MessageFlags.Ephemeral });
+});
 
-Also note that [core](https://github.com/discordjs/discord.js/tree/main/packages/core) supports an
-abstract `gateway` property that can be easily implemented, making this pretty comfortable to
-use in conjunction. Refer to the [Gateway documentation](https://discord.js.org/docs/packages/core/main/Gateway:Interface).
+await gateway.init('work_balancing_group', [GatewayDispatchEvents.InteractionCreate]);
+```
 
 ## Links
 
diff --git a/packages/redis-gateway/package.json b/packages/redis-gateway/package.json
@@ -53,6 +53,7 @@
 		"undici": "^5.22.0"
 	},
 	"devDependencies": {
+		"@discordjs/core": "workspace:^",
 		"@types/node": "16.18.25",
 		"cross-env": "^7.0.3",
 		"eslint": "^8.39.0",
diff --git a/packages/redis-gateway/src/discordEvents.ts b/packages/redis-gateway/src/discordEvents.ts
diff --git a/packages/redis-gateway/src/index.ts b/packages/redis-gateway/src/index.ts
diff --git a/yarn.lock b/yarn.lock
