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,