fix goroutine leak

This commit is contained in:
isther 2024-01-04 11:08:09 +08:00
parent c60aac2f18
commit 9fbb411eb6
1 changed files with 10 additions and 1 deletions

View File

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