From 412cc2e71a5b578c03f4c31108cd0e00f667e078 Mon Sep 17 00:00:00 2001
From: topi314 <git@topi.wtf>
Date: Sat, 16 Mar 2024 12:06:50 +0100
Subject: [PATCH] allow multiple gateway closes without blocking forever

---
 gateway/gateway_impl.go | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/gateway/gateway_impl.go b/gateway/gateway_impl.go
index f644bf88..79a31b65 100644
--- a/gateway/gateway_impl.go
+++ b/gateway/gateway_impl.go
@@ -42,10 +42,10 @@ type gatewayImpl struct {
 	closeHandlerFunc CloseHandlerFunc
 	token            string
 
-	conn          *websocket.Conn
-	connMu        sync.Mutex
-	heartbeatChan chan struct{}
-	status        Status
+	conn            *websocket.Conn
+	connMu          sync.Mutex
+	heartbeatCancel context.CancelFunc
+	status          Status
 
 	heartbeatInterval     time.Duration
 	lastHeartbeatSent     time.Time
@@ -132,9 +132,9 @@ func (g *gatewayImpl) Close(ctx context.Context) {
 }
 
 func (g *gatewayImpl) CloseWithCode(ctx context.Context, code int, message string) {
-	if g.heartbeatChan != nil {
+	if g.heartbeatCancel != nil {
 		g.config.Logger.Debug("closing heartbeat goroutines...")
-		g.heartbeatChan <- struct{}{}
+		g.heartbeatCancel()
 	}
 
 	g.connMu.Lock()
@@ -234,16 +234,16 @@ func (g *gatewayImpl) reconnect() {
 }
 
 func (g *gatewayImpl) heartbeat() {
-	if g.heartbeatChan == nil {
-		g.heartbeatChan = make(chan struct{})
-	}
+	ctx, cancel := context.WithCancel(context.Background())
+	g.heartbeatCancel = cancel
+
 	heartbeatTicker := time.NewTicker(g.heartbeatInterval)
 	defer heartbeatTicker.Stop()
 	defer g.config.Logger.Debug("exiting heartbeat goroutine")
 
 	for {
 		select {
-		case <-g.heartbeatChan:
+		case <-ctx.Done():
 			return
 
 		case <-heartbeatTicker.C: