feat: ctx + ping

example.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"log"
 	"time"
 
@@ -15,6 +16,9 @@ const (
 
 func main() {
 	ws := websocket.NewWebSocketClient(
+		// ctx
+		context.Background(),
+
 		// onOpen
 		func(ws *websocket.WebSocketClient) {
 			log.Println("Connected")

websocket/client.go

@@ -25,11 +25,12 @@ type WebSocketClient struct {
 }
 
 func NewWebSocketClient(
+	ctx context.Context,
 	onOpen func(ws *WebSocketClient),
 	onClose func(ws *WebSocketClient),
 	onMessage func(ws *WebSocketClient, payload []byte),
 ) *WebSocketClient {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(ctx)
 	return &WebSocketClient{
 		ctx:       ctx,
 		cancel:    cancel,
@@ -84,6 +85,7 @@ func (ws *WebSocketClient) Connect(
 				var payload []byte
 				_, payload, err := ws.conn.Read(ws.ctx)
 				if err != nil {
+					ws.cancel()
 					if ws.onClose != nil {
 						ws.onClose(ws)
 					}
@@ -94,6 +96,21 @@ func (ws *WebSocketClient) Connect(
 		}
 	}()
 
+	// Ping Sender
+	ws.wg.Add(1)
+	go func() {
+		defer ws.wg.Done()
+		pingTicker := time.NewTicker(time.Second * 10)
+		for {
+			select {
+			case <-ws.ctx.Done():
+				return
+			case <-pingTicker.C:
+				ws.conn.Ping(ws.ctx)
+			}
+		}
+	}()
+
 	return nil
 }
 
@@ -113,3 +130,7 @@ func (ws *WebSocketClient) Disconnect() {
 func (ws *WebSocketClient) SendJSON(v any) error {
 	return wsjson.Write(ws.ctx, ws.conn, v)
 }
+
+func (ws *WebSocketClient) SendPing() error {
+	return ws.conn.Ping(ws.ctx)
+}