From 54565f6af86f2cb47239c98a10ea00b000d7c355 Mon Sep 17 00:00:00 2001
From: minish <minish@riseup.net>
Date: Tue, 30 Dec 2025 03:23:47 -0500
Subject: [PATCH] feat: ctx + ping

---
 example.go          |  4 ++++
 websocket/client.go | 23 ++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/example.go b/example.go
index d1bc9dc..b6ff2c0 100644
--- a/example.go
+++ b/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")
diff --git a/websocket/client.go b/websocket/client.go
index 8282a32..291834b 100644
--- a/websocket/client.go
+++ b/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)
+}