feat(ws): metrics (#9005)

* feat(WebSocketManager): fetch status

* feat(WebSocketShard): heartbeat event

* chore: ci

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>

packages/ws/src/strategies/sharding/IShardingStrategy.ts

@@ -1,6 +1,7 @@
+import type { Collection } from '@discordjs/collection';
 import type { Awaitable } from '@discordjs/util';
 import type { GatewaySendPayload } from 'discord-api-types/v10';
-import type { WebSocketShardDestroyOptions } from '../../ws/WebSocketShard';
+import type { WebSocketShardDestroyOptions, WebSocketShardStatus } from '../../ws/WebSocketShard';
 
 /**
  * Strategies responsible for spawning, initializing connections, destroying shards, and relaying events
@@ -14,6 +15,10 @@ export interface IShardingStrategy {
 	 * Destroys all the shards
 	 */
 	destroy(options?: Omit<WebSocketShardDestroyOptions, 'recover'>): Awaitable<void>;
+	/**
+	 * Fetches the status of all the shards
+	 */
+	fetchStatus(): Awaitable<Collection<number, WebSocketShardStatus>>;
 	/**
 	 * Sends a payload to a shard
 	 */

packages/ws/src/strategies/sharding/SimpleShardingStrategy.ts

@@ -70,4 +70,11 @@ export class SimpleShardingStrategy implements IShardingStrategy {
 		if (!shard) throw new Error(`Shard ${shardId} not found`);
 		return shard.send(payload);
 	}
+
+	/**
+	 * {@inheritDoc IShardingStrategy.fetchStatus}
+	 */
+	public async fetchStatus() {
+		return this.shards.mapValues((shard) => shard.status);
+	}
 }

packages/ws/src/strategies/sharding/WorkerShardingStrategy.ts

@@ -5,7 +5,7 @@ import { Collection } from '@discordjs/collection';
 import type { GatewaySendPayload } from 'discord-api-types/v10';
 import { IdentifyThrottler } from '../../utils/IdentifyThrottler.js';
 import type { SessionInfo, WebSocketManager } from '../../ws/WebSocketManager';
-import type { WebSocketShardDestroyOptions, WebSocketShardEvents } from '../../ws/WebSocketShard';
+import type { WebSocketShardDestroyOptions, WebSocketShardEvents, WebSocketShardStatus } from '../../ws/WebSocketShard';
 import { managerToFetchingStrategyOptions, type FetchingStrategyOptions } from '../context/IContextFetchingStrategy.js';
 import type { IShardingStrategy } from './IShardingStrategy.js';
 
@@ -19,9 +19,11 @@ export enum WorkerSendPayloadOp {
 	Send,
 	SessionInfoResponse,
 	ShardCanIdentify,
+	FetchStatus,
 }
 
 export type WorkerSendPayload =
+	| { 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 }
@@ -35,11 +37,13 @@ export enum WorkerRecievePayloadOp {
 	RetrieveSessionInfo,
 	UpdateSessionInfo,
 	WaitForIdentify,
+	FetchStatusResponse,
 }
 
 export type WorkerRecievePayload =
 	// 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: WorkerRecievePayloadOp.Event; shardId: number }
+	| { nonce: number; op: WorkerRecievePayloadOp.FetchStatusResponse; status: WebSocketShardStatus }
 	| { nonce: number; op: WorkerRecievePayloadOp.RetrieveSessionInfo; shardId: number }
 	| { nonce: number; op: WorkerRecievePayloadOp.WaitForIdentify }
 	| { op: WorkerRecievePayloadOp.Connected; shardId: number }
@@ -72,6 +76,8 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 
 	private readonly destroyPromises = new Collection<number, () => void>();
 
+	private readonly fetchStatusPromises = new Collection<number, (status: WebSocketShardStatus) => void>();
+
 	private readonly throttler: IdentifyThrottler;
 
 	public constructor(manager: WebSocketManager, options: WorkerShardingStrategyOptions) {
@@ -179,18 +185,41 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 		worker.postMessage(payload);
 	}
 
+	/**
+	 * {@inheritDoc IShardingStrategy.fetchStatus}
+	 */
+	public async fetchStatus() {
+		const statuses = new Collection<number, WebSocketShardStatus>();
+
+		for (const [shardId, worker] of this.#workerByShardId.entries()) {
+			const nonce = Math.random();
+			const payload = {
+				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));
+			worker.postMessage(payload);
+
+			const status = await promise;
+			statuses.set(shardId, status);
+		}
+
+		return statuses;
+	}
+
 	private async onMessage(worker: Worker, payload: WorkerRecievePayload) {
 		switch (payload.op) {
 			case WorkerRecievePayloadOp.Connected: {
-				const resolve = this.connectPromises.get(payload.shardId)!;
-				resolve();
+				this.connectPromises.get(payload.shardId)?.();
 				this.connectPromises.delete(payload.shardId);
 				break;
 			}
 
 			case WorkerRecievePayloadOp.Destroyed: {
-				const resolve = this.destroyPromises.get(payload.shardId)!;
-				resolve();
+				this.destroyPromises.get(payload.shardId)?.();
 				this.destroyPromises.delete(payload.shardId);
 				break;
 			}
@@ -225,6 +254,12 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 				worker.postMessage(response);
 				break;
 			}
+
+			case WorkerRecievePayloadOp.FetchStatusResponse: {
+				this.fetchStatusPromises.get(payload.nonce)?.(payload.status);
+				this.fetchStatusPromises.delete(payload.nonce);
+				break;
+			}
 		}
 	}
 }

packages/ws/src/strategies/sharding/worker.ts

@@ -97,5 +97,21 @@ parentPort!
 			case WorkerSendPayloadOp.ShardCanIdentify: {
 				break;
 			}
+
+			case WorkerSendPayloadOp.FetchStatus: {
+				const shard = shards.get(payload.shardId);
+				if (!shard) {
+					throw new Error(`Shard ${payload.shardId} does not exist`);
+				}
+
+				const response = {
+					op: WorkerRecievePayloadOp.FetchStatusResponse,
+					status: shard.status,
+					nonce: payload.nonce,
+				} satisfies WorkerRecievePayload;
+
+				parentPort!.postMessage(response);
+				break;
+			}
 		}
 	});