MM-23503: Fix race in websocket_client writer (#14160)

* MM-23503: Fix race in websocket_client writer

We create a separate writer goroutine where all the writes happen.

* Fixing the case of double close

* Incorporate review comments

* Use CAS

* Fix incorrect comment

* Check if client is closed in pingHandler too

Co-authored-by: mattermod <mattermod@users.noreply.github.com>

model/websocket_client.go

@@ -7,6 +7,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"net/http"
+	"sync/atomic"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -17,6 +18,18 @@ const (
 	PING_TIMEOUT_BUFFER_SECONDS = 5
 )
 
+type msgType int
+
+const (
+	msgTypeJSON msgType = iota + 1
+	msgTypePong
+)
+
+type writeMessage struct {
+	msgType msgType
+	data    interface{}
+}
+
 // WebSocketClient stores the necessary information required to
 // communicate with a WebSocket endpoint.
 type WebSocketClient struct {
@@ -30,8 +43,13 @@ type WebSocketClient struct {
 	EventChannel       chan *WebSocketEvent    // The channel used to receive various events pushed from the server. For example: typing, posted
 	ResponseChannel    chan *WebSocketResponse // The channel used to receive responses for requests made to the server
 	ListenError        *AppError               // A field that is set if there was an abnormal closure of the WebSocket connection
-	pingTimeoutTimer   *time.Timer
+	writeChan          chan writeMessage
-	closeChannel       chan struct{}
+
+	pingTimeoutTimer *time.Timer
+	quitPingWatchdog chan struct{}
+
+	quitWriterChan chan struct{}
+	closed         int32
 }
 
 // NewWebSocketClient constructs a new WebSocket client with convenience
@@ -49,21 +67,22 @@ func NewWebSocketClientWithDialer(dialer *websocket.Dialer, url, authToken strin
 	}
 
 	client := &WebSocketClient{
-		url,
+		Url:                url,
-		url + API_URL_SUFFIX,
+		ApiUrl:             url + API_URL_SUFFIX,
-		url + API_URL_SUFFIX + "/websocket",
+		ConnectUrl:         url + API_URL_SUFFIX + "/websocket",
-		conn,
+		Conn:               conn,
-		authToken,
+		AuthToken:          authToken,
-		1,
+		Sequence:           1,
-		make(chan bool, 1),
+		PingTimeoutChannel: make(chan bool, 1),
-		make(chan *WebSocketEvent, 100),
+		EventChannel:       make(chan *WebSocketEvent, 100),
-		make(chan *WebSocketResponse, 100),
+		ResponseChannel:    make(chan *WebSocketResponse, 100),
-		nil,
+		writeChan:          make(chan writeMessage),
-		nil,
+		quitPingWatchdog:   make(chan struct{}),
-		make(chan struct{}),
+		quitWriterChan:     make(chan struct{}),
 	}
 
 	client.configurePingHandling()
+	go client.writer()
 
 	client.SendMessage(WEBSOCKET_AUTHENTICATION_CHALLENGE, map[string]interface{}{"token": authToken})
 
@@ -82,18 +101,29 @@ func NewWebSocketClient4WithDialer(dialer *websocket.Dialer, url, authToken stri
 	return NewWebSocketClientWithDialer(dialer, url, authToken)
 }
 
+// Connect creates a websocket connection with the given ConnectUrl.
+// This is racy and error-prone should not be used. Use any of the New* functions to create a websocket.
 func (wsc *WebSocketClient) Connect() *AppError {
 	return wsc.ConnectWithDialer(websocket.DefaultDialer)
 }
 
+// ConnectWithDialer creates a websocket connection with the given ConnectUrl using the dialer.
+// This is racy and error-prone and should not be used. Use any of the New* functions to create a websocket.
 func (wsc *WebSocketClient) ConnectWithDialer(dialer *websocket.Dialer) *AppError {
 	var err error
 	wsc.Conn, _, err = dialer.Dial(wsc.ConnectUrl, nil)
 	if err != nil {
 		return NewAppError("Connect", "model.websocket_client.connect_fail.app_error", nil, err.Error(), http.StatusInternalServerError)
 	}
-
+	// Super racy and should not be done anyways.
+	// All of this needs to be redesigned for v6.
 	wsc.configurePingHandling()
+	// If it has been closed before, we just restart the writer.
+	if atomic.CompareAndSwapInt32(&wsc.closed, 1, 0) {
+		wsc.writeChan = make(chan writeMessage)
+		wsc.quitWriterChan = make(chan struct{})
+		go wsc.writer()
+	}
 
 	wsc.EventChannel = make(chan *WebSocketEvent, 100)
 	wsc.ResponseChannel = make(chan *WebSocketResponse, 100)
@@ -104,16 +134,38 @@ func (wsc *WebSocketClient) ConnectWithDialer(dialer *websocket.Dialer) *AppErro
 }
 
 func (wsc *WebSocketClient) Close() {
+	// CAS to 1 and proceed. Return if already 1.
+	if !atomic.CompareAndSwapInt32(&wsc.closed, 0, 1) {
+		return
+	}
+	wsc.quitWriterChan <- struct{}{}
+	close(wsc.writeChan)
 	wsc.Conn.Close()
 }
 
+func (wsc *WebSocketClient) writer() {
+	for {
+		select {
+		case msg := <-wsc.writeChan:
+			switch msg.msgType {
+			case msgTypeJSON:
+				wsc.Conn.WriteJSON(msg.data)
+			case msgTypePong:
+				wsc.Conn.WriteMessage(websocket.PongMessage, []byte{})
+			}
+		case <-wsc.quitWriterChan:
+			return
+		}
+	}
+}
+
 func (wsc *WebSocketClient) Listen() {
 	go func() {
 		defer func() {
 			wsc.Conn.Close()
 			close(wsc.EventChannel)
 			close(wsc.ResponseChannel)
-			close(wsc.closeChannel)
+			close(wsc.quitPingWatchdog)
 		}()
 
 		for {
@@ -153,8 +205,10 @@ func (wsc *WebSocketClient) SendMessage(action string, data map[string]interface
 	req.Data = data
 
 	wsc.Sequence++
-
+	wsc.writeChan <- writeMessage{
-	wsc.Conn.WriteJSON(req)
+		msgType: msgTypeJSON,
+		data:    req,
+	}
 }
 
 // UserTyping will push a user_typing event out to all connected users
@@ -189,12 +243,17 @@ func (wsc *WebSocketClient) configurePingHandling() {
 }
 
 func (wsc *WebSocketClient) pingHandler(appData string) error {
+	if atomic.LoadInt32(&wsc.closed) == 1 {
+		return nil
+	}
 	if !wsc.pingTimeoutTimer.Stop() {
 		<-wsc.pingTimeoutTimer.C
 	}
 
 	wsc.pingTimeoutTimer.Reset(time.Second * (60 + PING_TIMEOUT_BUFFER_SECONDS))
-	wsc.Conn.WriteMessage(websocket.PongMessage, []byte{})
+	wsc.writeChan <- writeMessage{
+		msgType: msgTypePong,
+	}
 	return nil
 }
 
@@ -202,6 +261,6 @@ func (wsc *WebSocketClient) pingWatchdog() {
 	select {
 	case <-wsc.pingTimeoutTimer.C:
 		wsc.PingTimeoutChannel <- true
-	case <-wsc.closeChannel:
+	case <-wsc.quitPingWatchdog:
 	}
 }

model/websocket_client_test.go

@@ -0,0 +1,58 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package model
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/require"
+)
+
+func dummyWebsocketHandler(t *testing.T) http.HandlerFunc {
+	return func(w http.ResponseWriter, req *http.Request) {
+		upgrader := &websocket.Upgrader{
+			ReadBufferSize:  1024,
+			WriteBufferSize: 1024,
+		}
+		conn, err := upgrader.Upgrade(w, req, nil)
+		var buf []byte
+		for {
+			_, buf, err = conn.ReadMessage()
+			if err != nil {
+				break
+			}
+			t.Logf("%s\n", buf)
+			err = conn.WriteMessage(websocket.PingMessage, []byte("ping"))
+			if err != nil {
+				break
+			}
+		}
+		if _, ok := err.(*websocket.CloseError); !ok {
+			require.NoError(t, err)
+		}
+	}
+}
+
+// TestWebSocketRace needs to be run with -race to verify that
+// there is no race.
+func TestWebSocketRace(t *testing.T) {
+	s := httptest.NewServer(dummyWebsocketHandler(t))
+	defer s.Close()
+
+	url := strings.Replace(s.URL, "http://", "ws://", 1)
+	cli, err := NewWebSocketClient4(url, "authToken")
+	require.Nil(t, err)
+
+	cli.Listen()
+
+	for i := 0; i < 10; i++ {
+		time.Sleep(500 * time.Millisecond)
+		cli.UserTyping("channel", "parentId")
+	}
+}