From b32750d5456efcf62913d13e9301bfabb41313f4 Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Sun, 17 Aug 2025 00:17:14 +0800
Subject: [PATCH] fix(websocket): websocket concurrent write during ping pong

---
 internal/net/gphttp/websocket/manager.go | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/internal/net/gphttp/websocket/manager.go b/internal/net/gphttp/websocket/manager.go
index b9bc27d1..305df03d 100644
--- a/internal/net/gphttp/websocket/manager.go
+++ b/internal/net/gphttp/websocket/manager.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -24,6 +25,8 @@ type Manager struct {
 	lastPingTime     atomic.Value
 	readCh           chan []byte
 	err              error
+
+	writeLock sync.Mutex
 }
 
 var defaultUpgrader = websocket.Upgrader{
@@ -97,6 +100,7 @@ func (cm *Manager) PeriodicWrite(interval time.Duration, getData func() (any, er
 			cm.Close()
 			return
 		}
+
 		if err := cm.WriteJSON(data, interval); err != nil {
 			cm.err = err
 			cm.Close()
@@ -143,6 +147,9 @@ func (cm *Manager) WriteData(typ int, data []byte, timeout time.Duration) error
 	case <-cm.ctx.Done():
 		return cm.err
 	default:
+		cm.writeLock.Lock()
+		defer cm.writeLock.Unlock()
+
 		if err := cm.conn.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
 			return err
 		}
@@ -183,8 +190,12 @@ func (cm *Manager) Close() {
 }
 
 func (cm *Manager) GracefulClose() {
+	cm.writeLock.Lock()
+	defer cm.writeLock.Unlock()
+
 	_ = cm.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
 	_ = cm.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+
 	cm.Close()
 }
 
@@ -227,12 +238,7 @@ func (cm *Manager) readRoutine() {
 
 			if typ == websocket.TextMessage && string(data) == "ping" {
 				cm.lastPingTime.Store(time.Now())
-				if err := cm.conn.SetWriteDeadline(time.Now().Add(cm.pongWriteTimeout)); err != nil {
-					cm.err = fmt.Errorf("failed to set write deadline: %w", err)
-					cm.Close()
-					return
-				}
-				if err := cm.conn.WriteMessage(websocket.TextMessage, []byte("pong")); err != nil {
+				if err := cm.WriteData(websocket.TextMessage, []byte("pong"), cm.pongWriteTimeout); err != nil {
 					cm.err = fmt.Errorf("failed to write pong message: %w", err)
 					cm.Close()
 					return