feat: adding ping & pong handler

v1/client/client.go

@@ -169,7 +169,6 @@ type SafeWebsocketClient struct {
 	doneMap        *safemap.SafeMap[string, chan struct{}]
 
 	writeChan chan Message
-	pongChan  chan error
 }
 
 func (wsClient *SafeWebsocketClient) connect() error {
@@ -196,7 +195,6 @@ func (wsClient *SafeWebsocketClient) connect() error {
 
 	if wsClient.headers != nil {
 		for k, v := range *wsClient.headers {
-			fmt.Println(k, v)
 			header.Set(k, v)
 		}
 	}
@@ -220,23 +218,30 @@ func (wsClient *SafeWebsocketClient) connect() error {
 	wsClient.conn = conn
 	wsClient.isConnected = true
 
+	conn.SetPingHandler(func(pingData string) error {
+		if err := conn.SetReadDeadline(time.Now().Add(readDeadline)); err != nil {
+			log.Printf("error on read deadline: %v\n", err)
+			return err
+		}
+
+		wsClient.writeChan <- Message{
+			MessageType: MessageTypePong,
+			Data:        []byte(pingData),
+		}
+		return nil
+	})
+
+	conn.SetPongHandler(func(pingData string) error {
+		if err := conn.SetReadDeadline(time.Now().Add(readDeadline)); err != nil {
+			log.Printf("error on read deadline: %v\n", err)
+			return err
+		}
+		return nil
+	})
+
 	go wsClient.writePump()
 	go wsClient.readPump()
 
-	// conn.SetPingHandler(func(pingData string) error {
-	// 	wsClient.writeChan <- Message{
-	// 		MessageType: MessageTypePong,
-	// 		Data:        []byte(pingData),
-	// 	}
-
-	// 	select {
-	// 	case err := <-wsClient.pongChan:
-	// 		return err
-	// 	default:
-	// 	}
-	// 	return nil
-	// })
-
 	return nil
 }
 
@@ -265,9 +270,6 @@ func (wsClient *SafeWebsocketClient) writePump() {
 
 			if err := c.WriteMessage(int(data.MessageType), data.Data); err != nil {
 				log.Printf("error on write message: %v\n", err)
-				if data.MessageType == MessageTypePong {
-					wsClient.pongChan <- err
-				}
 				return
 			}
 		}
@@ -450,7 +452,7 @@ func (wsClient *SafeWebsocketClient) Close() error {
 		wsClient.conn.Close()
 	}
 	wsClient.isConnected = false
-	// close(wsClient.dataChannel)
+	close(wsClient.dataChannel)
 
 	return nil
 }