From b7a7d8e7b6f0986ffd6f1466909f31db2e39d268 Mon Sep 17 00:00:00 2001 From: Agniva De Sarker Date: Tue, 16 Nov 2021 19:43:43 +0530 Subject: [PATCH] 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 | 5 +++++ api4/websocket_test.go | 39 +++++++++++++++++++++++++++++++++++++++ app/web_conn.go | 25 ++++++++++++++++++++----- app/web_hub.go | 30 ++++-------------------------- app/web_hub_test.go | 4 ---- app/websocket_router.go | 2 -- model/websocket_client.go | 21 +++++++++++++++++++-- 7 files changed, 87 insertions(+), 39 deletions(-) diff --git a/api4/apitestlib.go b/api4/apitestlib.go index 4ef51b647c..4003c17899 100644 --- a/api4/apitestlib.go +++ b/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) } diff --git a/api4/websocket_test.go b/api4/websocket_test.go index 19b3ac7fd1..5f985441d7 100644 --- a/api4/websocket_test.go +++ b/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() diff --git a/app/web_conn.go b/app/web_conn.go index 2c30f50387..9f9fd1d2a8 100644 --- a/app/web_conn.go +++ b/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 - // to the cluster if it does not have it. + // This does not handle reconnect requests across nodes in a cluster. + // 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 } diff --git a/app/web_hub.go b/app/web_hub.go index 94d8c8b0dc..db96fa1286 100644 --- a/app/web_hub.go +++ b/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 { - // The hello message should only be sent when the conn wasn't found. + if webConn.IsAuthenticated() && webConn.reuseCount == 0 { + // 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 { diff --git a/app/web_hub_test.go b/app/web_hub_test.go index 6866a7fd6f..630a10ce33 100644 --- a/app/web_hub_test.go +++ b/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")) diff --git a/app/websocket_router.go b/app/websocket_router.go index 97e9cd9547..82884a17da 100644 --- a/app/websocket_router.go +++ b/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() { diff --git a/model/websocket_client.go b/model/websocket_client.go index 2bd8a3b7cc..df35695a8f 100644 --- a/model/websocket_client.go +++ b/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,