From 9fbb411eb67dd34d2ec24a851e4fca01ab3d2e89 Mon Sep 17 00:00:00 2001
From: isther <catisther@gmail.com>
Date: Thu, 4 Jan 2024 11:08:09 +0800
Subject: [PATCH] fix goroutine leak

---
 bybit_websocket.go | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/bybit_websocket.go b/bybit_websocket.go
index c7523e9..be19463 100644
--- a/bybit_websocket.go
+++ b/bybit_websocket.go
@@ -1,14 +1,16 @@
 package bybit_connector
 
 import (
+	"context"
 	"crypto/hmac"
 	"crypto/sha256"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
+	"time"
+
 	"github.com/google/uuid"
 	"github.com/gorilla/websocket"
-	"time"
 )
 
 type MessageHandler func(message string) error
@@ -43,6 +45,8 @@ type WebSocket struct {
 	maxAliveTime string
 	pingInterval int
 	onMessage    MessageHandler
+	ctx          context.Context
+	cancel       context.CancelFunc
 }
 
 type WebsocketOption func(*WebSocket)
@@ -111,6 +115,7 @@ func (b *WebSocket) Connect(args []string) error {
 
 	go b.handleIncomingMessages()
 
+	b.ctx, b.cancel = context.WithCancel(context.Background())
 	Ping(b)
 
 	return b.sendSubscription(args)
@@ -126,12 +131,16 @@ func Ping(b *WebSocket) {
 				if err := b.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
 					fmt.Println("Failed to send ping:", err)
 				}
+			case <-b.ctx.Done():
+				fmt.Println("Exit ping")
+				return
 			}
 		}
 	}()
 }
 
 func (b *WebSocket) Disconnect() error {
+	b.cancel()
 	return b.conn.Close()
 }