refactor: abstract identify throttling and correct max_concurrency handling (#9375)

* refactor: properly support max_concurrency ratelimit keys

* fix: properly block for same key

* chore: export session state

* chore: throttler no longer requires manager

* refactor: abstract throttlers

* chore: proper member order

* chore: remove leftover debug log

* chore: use @link tag in doc comment

Co-authored-by: Jiralite <33201955+Jiralite@users.noreply.github.com>

* chore: suggested changes

* fix(WebSocketShard): cancel identify if the shard closed in the meantime

* refactor(throttlers): support abort signals

* fix: memory leak

* chore: remove leftover

---------

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

packages/ws/src/throttling/IIdentifyThrottler.ts

@@ -0,0 +1,11 @@
+/**
+ * IdentifyThrottlers are responsible for dictating when a shard is allowed to identify.
+ *
+ * @see {@link https://discord.com/developers/docs/topics/gateway#sharding-max-concurrency}
+ */
+export interface IIdentifyThrottler {
+	/**
+	 * Resolves once the given shard should be allowed to identify, or rejects if the operation was aborted.
+	 */
+	waitForIdentify(shardId: number, signal: AbortSignal): Promise<void>;
+}

packages/ws/src/throttling/SimpleIdentifyThrottler.ts

@@ -0,0 +1,50 @@
+import { setTimeout as sleep } from 'node:timers/promises';
+import { Collection } from '@discordjs/collection';
+import { AsyncQueue } from '@sapphire/async-queue';
+import type { IIdentifyThrottler } from './IIdentifyThrottler';
+
+/**
+ * The state of a rate limit key's identify queue.
+ */
+export interface IdentifyState {
+	queue: AsyncQueue;
+	resetsAt: number;
+}
+
+/**
+ * Local, in-memory identify throttler.
+ */
+export class SimpleIdentifyThrottler implements IIdentifyThrottler {
+	private readonly states = new Collection<number, IdentifyState>();
+
+	public constructor(private readonly maxConcurrency: number) {}
+
+	/**
+	 * {@inheritDoc IIdentifyThrottler.waitForIdentify}
+	 */
+	public async waitForIdentify(shardId: number, signal: AbortSignal): Promise<void> {
+		const key = shardId % this.maxConcurrency;
+
+		const state = this.states.ensure(key, () => {
+			return {
+				queue: new AsyncQueue(),
+				resetsAt: Number.POSITIVE_INFINITY,
+			};
+		});
+
+		await state.queue.wait({ signal });
+
+		try {
+			const diff = state.resetsAt - Date.now();
+			if (diff <= 5_000) {
+				// To account for the latency the IDENTIFY payload goes through, we add a bit more wait time
+				const time = diff + Math.random() * 1_500;
+				await sleep(time);
+			}
+
+			state.resetsAt = Date.now() + 5_000;
+		} finally {
+			state.queue.shift();
+		}
+	}
+}