From 98420826bc2296fc392f17e8254cf4ad743ff5af Mon Sep 17 00:00:00 2001
From: DD <didinele.dev@gmail.com>
Date: Sat, 18 Mar 2023 21:32:50 +0200
Subject: [PATCH] fix(WebSocketShard): cancel initial heartbeat in destroy
 (#9244)

* fix(WebSocketShard): cancel initial heartbeat in destroy

* refactor: use try/catch/finally

* chore: add debug log
---
 packages/ws/src/ws/WebSocketShard.ts | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/packages/ws/src/ws/WebSocketShard.ts b/packages/ws/src/ws/WebSocketShard.ts
index a5cf43040..f7cf67080 100644
--- a/packages/ws/src/ws/WebSocketShard.ts
+++ b/packages/ws/src/ws/WebSocketShard.ts
@@ -95,6 +95,8 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {
 
 	private sendRateLimitState: SendRateLimitState = getInitialSendRateLimitState();
 
+	private initialHeartbeatTimeoutController: AbortController | null = null;
+
 	private heartbeatInterval: NodeJS.Timer | null = null;
 
 	private lastHeartbeatAt = -1;
@@ -203,6 +205,11 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {
 			clearInterval(this.heartbeatInterval);
 		}
 
+		if (this.initialHeartbeatTimeoutController) {
+			this.initialHeartbeatTimeoutController.abort();
+			this.initialHeartbeatTimeoutController = null;
+		}
+
 		this.lastHeartbeatAt = -1;
 
 		// Clear session state if applicable
@@ -568,7 +575,17 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {
 				const firstWait = Math.floor(payload.d.heartbeat_interval * jitter);
 				this.debug([`Preparing first heartbeat of the connection with a jitter of ${jitter}; waiting ${firstWait}ms`]);
 
-				await sleep(firstWait);
+				try {
+					const controller = new AbortController();
+					this.initialHeartbeatTimeoutController = controller;
+					await sleep(firstWait, undefined, { signal: controller.signal });
+				} catch {
+					this.debug(['Cancelled initial heartbeat due to #destroy being called']);
+					return;
+				} finally {
+					this.initialHeartbeatTimeoutController = null;
+				}
+
 				await this.heartbeat();
 
 				this.debug([`First heartbeat sent, starting to beat every ${payload.d.heartbeat_interval}ms`]);