Fix i/o timeout in WebSocket writes by adjusting WriteDeadline and co…

…ntrol message handling (#194) * fix: Prevent i/o timeout by setting WriteDeadline on each write * chore: Update gorilla/websocket to v1.5.3
hirokisan · Nov 4, 2024 · 0067576 · 0067576
1 parent 28847ad
commit 0067576
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 12 deletions.
diff --git a/go.mod b/go.mod
@@ -5,7 +5,7 @@ go 1.21
 require (
 	github.com/google/go-querystring v1.1.0
 	github.com/google/uuid v1.6.0
-	github.com/gorilla/websocket v1.5.0
+	github.com/gorilla/websocket v1.5.3
 	github.com/stretchr/testify v1.8.2
 )
 

diff --git a/go.sum b/go.sum
@@ -7,8 +7,8 @@ github.com/google/go-querystring v1.1.0 h1:AnCroh3fv4ZBgVIf1Iwtovgjaw/GiKJo8M8yD
 github.com/google/go-querystring v1.1.0/go.mod h1:Kcdr2DB4koayq7X8pmAG4sNG59So17icRSOU623lUBU=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
-github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

diff --git a/v5_ws_private.go b/v5_ws_private.go
@@ -131,8 +131,8 @@ func (s *V5WebsocketPrivateService) Start(ctx context.Context, errHandler ErrHan
 	go func() {
 		defer close(done)
 		defer s.connection.Close()
+
 		_ = s.connection.SetReadDeadline(time.Now().Add(60 * time.Second))
-		_ = s.connection.SetWriteDeadline(time.Now().Add(60 * time.Second))
 		s.connection.SetPongHandler(func(string) error {
 			_ = s.connection.SetReadDeadline(time.Now().Add(60 * time.Second))
 			return nil
@@ -251,7 +251,7 @@ func (s *V5WebsocketPrivateService) Run() error {
 func (s *V5WebsocketPrivateService) Ping() error {
 	// NOTE: It appears that two messages need to be sent.
 	// REF: https://github.com/hirokisan/bybit/pull/127#issuecomment-1537479346
-	if err := s.writeMessage(websocket.PingMessage, nil); err != nil {
+	if err := s.writeControl(websocket.PingMessage, nil); err != nil {
 		return err
 	}
 	if err := s.writeMessage(websocket.TextMessage, []byte(`{"op":"ping"}`)); err != nil {
@@ -262,7 +262,7 @@ func (s *V5WebsocketPrivateService) Ping() error {
 
 // Close :
 func (s *V5WebsocketPrivateService) Close() error {
-	if err := s.writeMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")); err != nil && !errors.Is(err, websocket.ErrCloseSent) {
+	if err := s.writeControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")); err != nil && !errors.Is(err, websocket.ErrCloseSent) {
 		return err
 	}
 	return nil
@@ -272,8 +272,19 @@ func (s *V5WebsocketPrivateService) writeMessage(messageType int, body []byte) e
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	_ = s.connection.SetWriteDeadline(time.Now().Add(60 * time.Second))
 	if err := s.connection.WriteMessage(messageType, body); err != nil {
 		return err
 	}
 	return nil
 }
+
+func (s *V5WebsocketPrivateService) writeControl(messageType int, body []byte) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if err := s.connection.WriteControl(messageType, body, time.Now().Add(60*time.Second)); err != nil {
+		return err
+	}
+	return nil
+}
diff --git a/v5_ws_public.go b/v5_ws_public.go
@@ -148,7 +148,6 @@ func (s *V5WebsocketPublicService) Start(ctx context.Context, errHandler ErrHand
 		defer s.connection.Close()
 
 		_ = s.connection.SetReadDeadline(time.Now().Add(60 * time.Second))
-		_ = s.connection.SetWriteDeadline(time.Now().Add(60 * time.Second))
 		s.connection.SetPongHandler(func(string) error {
 			_ = s.connection.SetReadDeadline(time.Now().Add(60 * time.Second))
 			return nil
@@ -283,7 +282,7 @@ func (s *V5WebsocketPublicService) Run() error {
 func (s *V5WebsocketPublicService) Ping() error {
 	// NOTE: It appears that two messages need to be sent.
 	// REF: https://github.com/hirokisan/bybit/pull/127#issuecomment-1537479346
-	if err := s.writeMessage(websocket.PingMessage, nil); err != nil {
+	if err := s.writeControl(websocket.PingMessage, nil); err != nil {
 		return err
 	}
 	if err := s.writeMessage(websocket.TextMessage, []byte(`{"op":"ping"}`)); err != nil {
@@ -294,7 +293,7 @@ func (s *V5WebsocketPublicService) Ping() error {
 
 // Close :
 func (s *V5WebsocketPublicService) Close() error {
-	if err := s.writeMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")); err != nil && !errors.Is(err, websocket.ErrCloseSent) {
+	if err := s.writeControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")); err != nil && !errors.Is(err, websocket.ErrCloseSent) {
 		return err
 	}
 	return nil
@@ -304,8 +303,19 @@ func (s *V5WebsocketPublicService) writeMessage(messageType int, body []byte) er
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	_ = s.connection.SetWriteDeadline(time.Now().Add(60 * time.Second))
 	if err := s.connection.WriteMessage(messageType, body); err != nil {
 		return err
 	}
 	return nil
 }
+
+func (s *V5WebsocketPublicService) writeControl(messageType int, body []byte) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if err := s.connection.WriteControl(messageType, body, time.Now().Add(60*time.Second)); err != nil {
+		return err
+	}
+	return nil
+}
diff --git a/v5_ws_trade.go b/v5_ws_trade.go
@@ -83,8 +83,8 @@ func (s *V5WebsocketTradeService) Start(ctx context.Context, errHandler ErrHandl
 	go func() {
 		defer close(done)
 		defer s.connection.Close()
+
 		_ = s.connection.SetReadDeadline(time.Now().Add(60 * time.Second))
-		_ = s.connection.SetWriteDeadline(time.Now().Add(60 * time.Second))
 		s.connection.SetPongHandler(func(string) error {
 			_ = s.connection.SetReadDeadline(time.Now().Add(60 * time.Second))
 			return nil
@@ -154,7 +154,7 @@ func (s *V5WebsocketTradeService) Run() error {
 func (s *V5WebsocketTradeService) Ping() error {
 	// NOTE: It appears that two messages need to be sent.
 	// REF: https://github.com/hirokisan/bybit/pull/127#issuecomment-1537479346
-	if err := s.writeMessage(websocket.PingMessage, nil); err != nil {
+	if err := s.writeControl(websocket.PingMessage, nil); err != nil {
 		return err
 	}
 	if err := s.writeMessage(websocket.TextMessage, []byte(`{"op":"ping"}`)); err != nil {
@@ -165,7 +165,7 @@ func (s *V5WebsocketTradeService) Ping() error {
 
 // Close :
 func (s *V5WebsocketTradeService) Close() error {
-	if err := s.writeMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")); err != nil && !errors.Is(err, websocket.ErrCloseSent) {
+	if err := s.writeControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")); err != nil && !errors.Is(err, websocket.ErrCloseSent) {
 		return err
 	}
 	return nil
@@ -175,8 +175,19 @@ func (s *V5WebsocketTradeService) writeMessage(messageType int, body []byte) err
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	_ = s.connection.SetWriteDeadline(time.Now().Add(60 * time.Second))
 	if err := s.connection.WriteMessage(messageType, body); err != nil {
 		return err
 	}
 	return nil
 }
+
+func (s *V5WebsocketTradeService) writeControl(messageType int, body []byte) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if err := s.connection.WriteControl(messageType, body, time.Now().Add(60*time.Second)); err != nil {
+		return err
+	}
+	return nil
+}