MM-39612: Make acquiring and removing connections atomic (#18982)

* MM-39612: Make acquiring and removing connections atomic

The reconnect phase of a websocket was split into two parts:
one where we check if a connection with a given connectionID
exists or not. And second, where we remove that connection
and insert the new connection again in the index.

This would lead to a race where it would be possible
for 2 concurrent requests for the same connectionID to go through
which would lead to separate goroutines working on the same dead queue.

We simplify this by removing the connection from the index
in the check connection stage itself. And then just add that
during register phase.

And to distinguish between a fresh and an old connection, we add
a new field called reuseCount.

While here, we also cleanup some old comments and add more
in some places.

https://mattermost.atlassian.net/browse/MM-39612

```release-note
NONE
```

* remove unused method

```release-note
NONE
```

* race test

```release-note
NONE
```

api4/apitestlib.go

@@ -18,6 +18,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/gorilla/websocket"
 	s3 "github.com/minio/minio-go/v7"
 	"github.com/minio/minio-go/v7/pkg/credentials"
 	"github.com/stretchr/testify/require"
@@ -484,6 +485,10 @@ func (th *TestHelper) CreateWebSocketClient() (*model.WebSocketClient, error) {
 	return model.NewWebSocketClient4(fmt.Sprintf("ws://localhost:%v", th.App.Srv().ListenAddr.Port), th.Client.AuthToken)
 }
 
+func (th *TestHelper) CreateReliableWebSocketClient(connID string, seqNo int) (*model.WebSocketClient, error) {
+	return model.NewReliableWebSocketClientWithDialer(websocket.DefaultDialer, fmt.Sprintf("ws://localhost:%v", th.App.Srv().ListenAddr.Port), th.Client.AuthToken, connID, seqNo, true)
+}
+
 func (th *TestHelper) CreateWebSocketSystemAdminClient() (*model.WebSocketClient, error) {
 	return model.NewWebSocketClient4(fmt.Sprintf("ws://localhost:%v", th.App.Srv().ListenAddr.Port), th.SystemAdminClient.AuthToken)
 }

api4/websocket_test.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"net/http"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -199,6 +200,44 @@ func TestWebsocketOriginSecurity(t *testing.T) {
 	th.App.UpdateConfig(func(cfg *model.Config) { *cfg.ServiceSettings.AllowCorsFrom = "" })
 }
 
+func TestWebSocketReconnectRace(t *testing.T) {
+	th := Setup(t).InitBasic()
+	defer th.TearDown()
+
+	th.App.UpdateConfig(func(cfg *model.Config) {
+		*cfg.ServiceSettings.EnableReliableWebSockets = true
+	})
+
+	WebSocketClient, err := th.CreateWebSocketClient()
+	require.NoError(t, err)
+	defer WebSocketClient.Close()
+	WebSocketClient.Listen()
+
+	ev := <-WebSocketClient.EventChannel
+	require.Equal(t, model.WebsocketEventHello, ev.EventType())
+	evData := ev.GetData()
+	connID := evData["connection_id"].(string)
+	seq := int(ev.GetSequence())
+
+	var wg sync.WaitGroup
+	n := 10
+	wg.Add(n)
+
+	WebSocketClient.Close()
+
+	for i := 0; i < n; i++ {
+		go func() {
+			defer wg.Done()
+			ws, err := th.CreateReliableWebSocketClient(connID, seq+1)
+			require.NoError(t, err)
+			defer ws.Close()
+			ws.Listen()
+		}()
+	}
+
+	wg.Wait()
+}
+
 func TestWebSocketStatuses(t *testing.T) {
 	th := Setup(t).InitBasic()
 	defer th.TearDown()

app/web_conn.go

@@ -57,6 +57,7 @@ type WebConnConfig struct {
 	Locale       string
 	ConnectionID string
 	Active       bool
+	ReuseCount   int
 
 	// These aren't necessary to be exported to api layer.
 	sequence         int
@@ -94,7 +95,13 @@ type WebConn struct {
 	// to this webConn or not.
 	// It is not used as an atomic, because there is no need to.
 	// So do not use this outside the web hub.
-	active       bool
+	active bool
+	// reuseCount indicates how many times this connection has been reused.
+	// This is used to differentiate between a fresh connection and
+	// a reused connection.
+	// It's theoretically possible for this number to wrap around. But we
+	// leave that as an edge-case.
+	reuseCount   int
 	sessionToken atomic.Value
 	session      atomic.Value
 	connectionID atomic.Value
@@ -111,6 +118,7 @@ type CheckConnResult struct {
 	ActiveQueue      chan model.WebSocketMessage
 	DeadQueue        []*model.WebSocketEvent
 	DeadQueuePointer int
+	ReuseCount       int
 }
 
 // PopulateWebConnConfig checks if the connection id already exists in the hub,
@@ -120,8 +128,8 @@ func (a *App) PopulateWebConnConfig(s *model.Session, cfg *WebConnConfig, seqVal
 		return nil, fmt.Errorf("invalid connection id: %s", cfg.ConnectionID)
 	}
 
-	// TODO: the method should internally forward the request
+	// This does not handle reconnect requests across nodes in a cluster.
-	// to the cluster if it does not have it.
+	// It falls back to the non-reliable case in that scenario.
 	res := a.CheckWebConn(s.UserId, cfg.ConnectionID)
 	if res == nil {
 		// If the connection is not present, then we assume either timeout,
@@ -133,6 +141,7 @@ func (a *App) PopulateWebConnConfig(s *model.Session, cfg *WebConnConfig, seqVal
 		cfg.deadQueue = res.DeadQueue
 		cfg.deadQueuePointer = res.DeadQueuePointer
 		cfg.Active = false
+		cfg.ReuseCount = res.ReuseCount
 		// Now we get the sequence number
 		if seqVal == "" {
 			// Sequence_number must be sent with connection id.
@@ -188,6 +197,7 @@ func (a *App) NewWebConn(cfg *WebConnConfig) *WebConn {
 		T:                  cfg.TFunc,
 		Locale:             cfg.Locale,
 		active:             cfg.Active,
+		reuseCount:         cfg.ReuseCount,
 		endWritePump:       make(chan struct{}),
 		pumpFinished:       make(chan struct{}),
 		pluginPosted:       make(chan pluginWSPostedHook, 10),
@@ -525,12 +535,17 @@ func (wc *WebConn) addToDeadQueue(msg *model.WebSocketEvent) {
 // the latest element in the dead queue, which would mean there is no message loss.
 func (wc *WebConn) hasMsgLoss() bool {
 	var index int
+	// deadQueuePointer = 0 means either no msg written or the pointer
+	// has rolled over to its starting position.
 	if wc.deadQueuePointer == 0 {
+		// If last entry is nil, it means no msg is written.
 		if wc.deadQueue[deadQueueSize-1] == nil {
-			return false // No msg written
+			return false
 		}
+		// If it's not nil, that means it has rolled over to start, and we
+		// check the last position.
 		index = deadQueueSize - 1
-	} else {
+	} else { // deadQueuePointer != 0 means it's somewhere in the middle.
 		index = wc.deadQueuePointer - 1
 	}
 

app/web_hub.go

@@ -432,7 +432,7 @@ func (h *Hub) Start() {
 				webSessionMessage.isRegistered <- isRegistered
 			case req := <-h.checkConn:
 				var res *CheckConnResult
-				conn := connIndex.GetInactiveByConnectionID(req.userID, req.connectionID)
+				conn := connIndex.RemoveInactiveByConnectionID(req.userID, req.connectionID)
 				if conn != nil {
 					res = &CheckConnResult{
 						ConnectionID:     req.connectionID,
@@ -440,20 +440,13 @@ func (h *Hub) Start() {
 						ActiveQueue:      conn.send,
 						DeadQueue:        conn.deadQueue,
 						DeadQueuePointer: conn.deadQueuePointer,
+						ReuseCount:       conn.reuseCount + 1,
 					}
 				}
 				req.result <- res
 			case <-ticker.C:
 				connIndex.RemoveInactiveConnections()
 			case webConn := <-h.register:
-				var oldConn *WebConn
-				if *h.srv.Config().ServiceSettings.EnableReliableWebSockets {
-					// Delete the old conn from connIndex if it exists.
-					oldConn = connIndex.RemoveInactiveByConnectionID(
-						webConn.GetSession().UserId,
-						webConn.GetConnectionID())
-				}
-
 				// Mark the current one as active.
 				// There is no need to check if it was inactive or not,
 				// we will anyways need to make it active.
@@ -462,8 +455,8 @@ func (h *Hub) Start() {
 				connIndex.Add(webConn)
 				atomic.StoreInt64(&h.connectionCount, int64(connIndex.AllActive()))
 
-				if webConn.IsAuthenticated() && oldConn == nil {
+				if webConn.IsAuthenticated() && webConn.reuseCount == 0 {
-					// The hello message should only be sent when the conn wasn't found.
+					// The hello message should only be sent when the reuseCount is 0.
 					// i.e in server restart, or long timeout, or fresh connection case.
 					// In case of seq number not found in dead queue, it is handled by
 					// the webconn write pump.
@@ -660,21 +653,6 @@ func (i *hubConnectionIndex) All() map[*WebConn]int {
 	return i.byConnection
 }
 
-// GetInactiveByConnectionID returns an inactive connection for the given
-// userID and connectionID.
-func (i *hubConnectionIndex) GetInactiveByConnectionID(userID, connectionID string) *WebConn {
-	// To handle empty sessions.
-	if userID == "" {
-		return nil
-	}
-	for _, conn := range i.ForUser(userID) {
-		if conn.GetConnectionID() == connectionID && !conn.active {
-			return conn
-		}
-	}
-	return nil
-}
-
 // RemoveInactiveByConnectionID removes an inactive connection for the given
 // userID and connectionID.
 func (i *hubConnectionIndex) RemoveInactiveByConnectionID(userID, connectionID string) *WebConn {

app/web_hub_test.go

@@ -316,10 +316,6 @@ func TestHubConnIndexInactive(t *testing.T) {
 	connIndex.Add(wc2)
 	connIndex.Add(wc3)
 
-	assert.Nil(t, connIndex.GetInactiveByConnectionID(wc2.UserId, "conn2"))
-	assert.NotNil(t, connIndex.GetInactiveByConnectionID(wc2.UserId, "conn3"))
-	assert.Nil(t, connIndex.GetInactiveByConnectionID(wc1.UserId, "conn3"))
-
 	assert.Nil(t, connIndex.RemoveInactiveByConnectionID(wc2.UserId, "conn2"))
 	assert.NotNil(t, connIndex.RemoveInactiveByConnectionID(wc2.UserId, "conn3"))
 	assert.Nil(t, connIndex.RemoveInactiveByConnectionID(wc1.UserId, "conn3"))

app/websocket_router.go

@@ -56,8 +56,6 @@ func (wr *WebSocketRouter) ServeWebSocket(conn *WebConn, r *model.WebSocketReque
 		conn.SetSessionToken(session.Token)
 		conn.UserId = session.UserId
 
-		// TODO: Same logic to reconnect queue as api4/websocket.go
-
 		conn.App.HubRegister(conn)
 
 		conn.App.Srv().Go(func() {

model/websocket_client.go

@@ -6,6 +6,7 @@ package model
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"sync/atomic"
 	"time"
@@ -65,10 +66,26 @@ func NewWebSocketClient(url, authToken string) (*WebSocketClient, error) {
 	return NewWebSocketClientWithDialer(websocket.DefaultDialer, url, authToken)
 }
 
+func NewReliableWebSocketClientWithDialer(dialer *websocket.Dialer, url, authToken, connID string, seqNo int, withAuthHeader bool) (*WebSocketClient, error) {
+	connectURL := url + APIURLSuffix + "/websocket" + fmt.Sprintf("?connection_id=%s&sequence_number=%d", connID, seqNo)
+	var header http.Header
+	if withAuthHeader {
+		header = http.Header{
+			"Authorization": []string{"Bearer " + authToken},
+		}
+	}
+
+	return makeClient(dialer, url, connectURL, authToken, header)
+}
+
 // NewWebSocketClientWithDialer constructs a new WebSocket client with convenience
 // methods for talking to the server using a custom dialer.
 func NewWebSocketClientWithDialer(dialer *websocket.Dialer, url, authToken string) (*WebSocketClient, error) {
-	conn, _, err := dialer.Dial(url+APIURLSuffix+"/websocket", nil)
+	return makeClient(dialer, url, url+APIURLSuffix+"/websocket", authToken, nil)
+}
+
+func makeClient(dialer *websocket.Dialer, url, connectURL, authToken string, header http.Header) (*WebSocketClient, error) {
+	conn, _, err := dialer.Dial(connectURL, header)
 	if err != nil {
 		return nil, NewAppError("NewWebSocketClient", "model.websocket_client.connect_fail.app_error", nil, err.Error(), http.StatusInternalServerError)
 	}
@@ -76,7 +93,7 @@ func NewWebSocketClientWithDialer(dialer *websocket.Dialer, url, authToken strin
 	client := &WebSocketClient{
 		URL:                url,
 		APIURL:             url + APIURLSuffix,
-		ConnectURL:         url + APIURLSuffix + "/websocket",
+		ConnectURL:         connectURL,
 		Conn:               conn,
 		AuthToken:          authToken,
 		Sequence:           1,