diff --git a/api4/websocket.go b/api4/websocket.go index 2608b5ebbf..6697abd489 100644 --- a/api4/websocket.go +++ b/api4/websocket.go @@ -5,10 +5,10 @@ package api4 import ( "net/http" - "strconv" "github.com/gorilla/websocket" + "github.com/mattermost/mattermost-server/v5/app" "github.com/mattermost/mattermost-server/v5/model" "github.com/mattermost/mattermost-server/v5/shared/mlog" ) @@ -36,53 +36,34 @@ func connectWebSocket(c *Context, w http.ResponseWriter, r *http.Request) { return } - wc := c.App.NewWebConn(ws, *c.App.Session(), c.App.T, "") - - if *c.App.Config().ServiceSettings.EnableReliableWebSockets { - connID := r.URL.Query().Get(connectionIDParam) - if connID == "" { - // If not present, we assume client is not capable yet, or it's a fresh connection. - // We just create a new ID. - connID = model.NewId() - } else { - if !model.IsValidId(connID) { - mlog.Error("Invalid connection ID", mlog.String("id", connID)) - wc.WebSocket.Close() - return - } - // If present, we check if it's present in the connection manager. - // TODO: the connection manager internally should forward the request - // to the cluster if it does not have it. - // - // If the connection is not present, then we assume either timeout, - // or server restart. In that case, we set a new one. - // - // Now we get the sequence number - seqVal := r.URL.Query().Get(sequenceNumberParam) - if seqVal == "" { - // Sequence_number must be sent with connection id. - // A client must be either non-compliant or fully compliant. - mlog.Error("Sequence number not present in websocket request") - wc.WebSocket.Close() - return - } - seq, err := strconv.Atoi(seqVal) - if err != nil || seq < 0 { - mlog.Error("Invalid sequence number set in query param", - mlog.String("query", seqVal), - mlog.Err(err)) - wc.WebSocket.Close() - return - } - wc.Sequence = int64(seq) - // Now if there have been past entries to be back-filled, we do it. - // First we find the right sequence number point. - // We start consuming from dead queue first, and then move to active queue - } - // In case of fresh connection id, sequence number is already zero. - wc.SetConnectionID(connID) + // We initialize webconn with all the necessary data. + // If the queues are empty, they are initialized in the constructor. + cfg := &app.WebConnConfig{ + WebSocket: ws, + Session: *c.App.Session(), + TFunc: c.App.T, + Locale: "", + Active: true, } + if *c.App.Config().ServiceSettings.EnableReliableWebSockets { + cfg.ConnectionID = r.URL.Query().Get(connectionIDParam) + if cfg.ConnectionID == "" || c.App.Session().UserId == "" { + // If not present, we assume client is not capable yet, or it's a fresh connection. + // We just create a new ID. + cfg.ConnectionID = model.NewId() + // In case of fresh connection id, sequence number is already zero. + } else { + cfg, err = c.App.PopulateWebConnConfig(cfg, r.URL.Query().Get(sequenceNumberParam)) + if err != nil { + mlog.Warn("Error while populating webconn config", mlog.String("id", r.URL.Query().Get(connectionIDParam)), mlog.Err(err)) + ws.Close() + return + } + } + } + + wc := c.App.NewWebConn(cfg) if c.App.Session().UserId != "" { c.App.HubRegister(wc) } diff --git a/app/app_iface.go b/app/app_iface.go index be0f0243ee..8cc62b9259 100644 --- a/app/app_iface.go +++ b/app/app_iface.go @@ -18,7 +18,6 @@ import ( "time" "github.com/dyatlov/go-opengraph/opengraph" - "github.com/gorilla/websocket" "github.com/mattermost/mattermost-server/v5/audit" "github.com/mattermost/mattermost-server/v5/einterfaces" "github.com/mattermost/mattermost-server/v5/model" @@ -243,7 +242,7 @@ type AppIface interface { // function is only exposed to sysadmins and the possibility of this edge case is relatively small. MoveChannel(team *model.Team, channel *model.Channel, user *model.User) *model.AppError // NewWebConn returns a new WebConn instance. - NewWebConn(ws *websocket.Conn, session model.Session, t i18n.TranslateFunc, locale string) *WebConn + NewWebConn(cfg *WebConnConfig) *WebConn // NewWebHub creates a new Hub. NewWebHub() *Hub // NotifySessionsExpired is called periodically from the job server to notify any mobile sessions that have expired. @@ -261,6 +260,9 @@ type AppIface interface { DoActionRequest(rawURL string, body []byte) (*http.Response, *model.AppError) // PermanentDeleteBot permanently deletes a bot and its corresponding user. PermanentDeleteBot(botUserId string) *model.AppError + // PopulateWebConnConfig checks if the connection id already exists in the hub, + // and if so, accordingly populates the other fields of the webconn. + PopulateWebConnConfig(cfg *WebConnConfig, seqVal string) (*WebConnConfig, error) // PromoteGuestToUser Convert user's roles and all his mermbership's roles from // guest roles to regular user roles. PromoteGuestToUser(user *model.User, requestorId string) *model.AppError @@ -419,6 +421,7 @@ type AppIface interface { CheckUserPostflightAuthenticationCriteria(user *model.User) *model.AppError CheckUserPreflightAuthenticationCriteria(user *model.User, mfaToken string) *model.AppError CheckValidDomains(team *model.Team) *model.AppError + CheckWebConn(userID, connectionID string) *CheckConnResult ClearChannelMembersCache(channelID string) ClearSessionCacheForAllUsers() ClearSessionCacheForAllUsersSkipClusterSend() diff --git a/app/opentracing/opentracing_layer.go b/app/opentracing/opentracing_layer.go index 7dfdc58ebc..e09bf10d33 100644 --- a/app/opentracing/opentracing_layer.go +++ b/app/opentracing/opentracing_layer.go @@ -18,7 +18,6 @@ import ( "time" "github.com/dyatlov/go-opengraph/opengraph" - "github.com/gorilla/websocket" "github.com/mattermost/mattermost-server/v5/app" "github.com/mattermost/mattermost-server/v5/audit" "github.com/mattermost/mattermost-server/v5/einterfaces" @@ -1353,6 +1352,23 @@ func (a *OpenTracingAppLayer) CheckValidDomains(team *model.Team) *model.AppErro return resultVar0 } +func (a *OpenTracingAppLayer) CheckWebConn(userID string, connectionID string) *app.CheckConnResult { + origCtx := a.ctx + span, newCtx := tracing.StartSpanWithParentByContext(a.ctx, "app.CheckWebConn") + + a.ctx = newCtx + a.app.Srv().Store.SetContext(newCtx) + defer func() { + a.app.Srv().Store.SetContext(origCtx) + a.ctx = origCtx + }() + + defer span.Finish() + resultVar0 := a.app.CheckWebConn(userID, connectionID) + + return resultVar0 +} + func (a *OpenTracingAppLayer) ClearChannelMembersCache(channelID string) { origCtx := a.ctx span, newCtx := tracing.StartSpanWithParentByContext(a.ctx, "app.ClearChannelMembersCache") @@ -11477,7 +11493,7 @@ func (a *OpenTracingAppLayer) NewPluginAPI(manifest *model.Manifest) plugin.API return resultVar0 } -func (a *OpenTracingAppLayer) NewWebConn(ws *websocket.Conn, session model.Session, t i18n.TranslateFunc, locale string) *app.WebConn { +func (a *OpenTracingAppLayer) NewWebConn(cfg *app.WebConnConfig) *app.WebConn { origCtx := a.ctx span, newCtx := tracing.StartSpanWithParentByContext(a.ctx, "app.NewWebConn") @@ -11489,7 +11505,7 @@ func (a *OpenTracingAppLayer) NewWebConn(ws *websocket.Conn, session model.Sessi }() defer span.Finish() - resultVar0 := a.app.NewWebConn(ws, session, t, locale) + resultVar0 := a.app.NewWebConn(cfg) return resultVar0 } @@ -11988,6 +12004,28 @@ func (a *OpenTracingAppLayer) PluginContext() *plugin.Context { return resultVar0 } +func (a *OpenTracingAppLayer) PopulateWebConnConfig(cfg *app.WebConnConfig, seqVal string) (*app.WebConnConfig, error) { + origCtx := a.ctx + span, newCtx := tracing.StartSpanWithParentByContext(a.ctx, "app.PopulateWebConnConfig") + + a.ctx = newCtx + a.app.Srv().Store.SetContext(newCtx) + defer func() { + a.app.Srv().Store.SetContext(origCtx) + a.ctx = origCtx + }() + + defer span.Finish() + resultVar0, resultVar1 := a.app.PopulateWebConnConfig(cfg, seqVal) + + if resultVar1 != nil { + span.LogFields(spanlog.Error(resultVar1)) + ext.Error.Set(span, true) + } + + return resultVar0, resultVar1 +} + func (a *OpenTracingAppLayer) PostActionCookieSecret() []byte { origCtx := a.ctx span, newCtx := tracing.StartSpanWithParentByContext(a.ctx, "app.PostActionCookieSecret") diff --git a/app/web_conn.go b/app/web_conn.go index 1de8b76fe3..5ffb4e7069 100644 --- a/app/web_conn.go +++ b/app/web_conn.go @@ -6,9 +6,11 @@ package app import ( "bytes" "encoding/json" + "errors" "fmt" "net" "net/http" + "strconv" "sync" "sync/atomic" "time" @@ -32,6 +34,21 @@ const ( deadQueueSize = 128 // Approximated from /proc/sys/net/core/wmem_default / 2048 (avg msg size) ) +type WebConnConfig struct { + WebSocket *websocket.Conn + Session model.Session + TFunc i18n.TranslateFunc + Locale string + ConnectionID string + Active bool + + // These aren't necessary to be exported to api layer. + sequence int + activeQueue chan model.WebSocketMessage + deadQueue []*model.WebSocketEvent + deadQueuePointer int +} + // WebConn represents a single websocket connection to a user. // It contains all the necessary state to manage sending/receiving data to/from // a websocket. @@ -53,28 +70,80 @@ type WebConn struct { // It basically acts as the user-space socket buffer, and is used // to resuscitate any messages that might have got lost when the connection is broken. // It is implemented by using a circular buffer to keep it fast. - deadQueue []model.WebSocketMessage - deadQueuePointer int // Pointer which indicates the next slot to insert. - sessionToken atomic.Value - session atomic.Value - connectionID atomic.Value - endWritePump chan struct{} - pumpFinished chan struct{} + deadQueue []*model.WebSocketEvent + // Pointer which indicates the next slot to insert. + // It is only to be incremented during writing or clearing the queue. + deadQueuePointer int + // active indicates whether there is an open websocket connection attached + // 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 + sessionToken atomic.Value + session atomic.Value + connectionID atomic.Value + endWritePump chan struct{} + pumpFinished chan struct{} +} + +// CheckConnResult indicates whether a connectionID was present in the hub or not. +// And if so, contains the active and dead queue details. +type CheckConnResult struct { + ConnectionID string + UserID string + ActiveQueue chan model.WebSocketMessage + DeadQueue []*model.WebSocketEvent + DeadQueuePointer int +} + +// PopulateWebConnConfig checks if the connection id already exists in the hub, +// and if so, accordingly populates the other fields of the webconn. +func (a *App) PopulateWebConnConfig(cfg *WebConnConfig, seqVal string) (*WebConnConfig, error) { + if !model.IsValidId(cfg.ConnectionID) { + 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. + res := a.CheckWebConn(a.Session().UserId, cfg.ConnectionID) + if res == nil { + // If the connection is not present, then we assume either timeout, + // or server restart. In that case, we set a new one. + cfg.ConnectionID = model.NewId() + } else { + // Connection is present, we get the active queue, dead queue + cfg.activeQueue = res.ActiveQueue + cfg.deadQueue = res.DeadQueue + cfg.deadQueuePointer = res.DeadQueuePointer + cfg.Active = false + // Now we get the sequence number + if seqVal == "" { + // Sequence_number must be sent with connection id. + // A client must be either non-compliant or fully compliant. + return nil, errors.New("Sequence number not present in websocket request") + } + var err error + cfg.sequence, err = strconv.Atoi(seqVal) + if err != nil || cfg.sequence < 0 { + return nil, fmt.Errorf("invalid sequence number %s in query param: %v", seqVal, err) + } + } + return cfg, nil } // NewWebConn returns a new WebConn instance. -func (a *App) NewWebConn(ws *websocket.Conn, session model.Session, t i18n.TranslateFunc, locale string) *WebConn { - if session.UserId != "" { +func (a *App) NewWebConn(cfg *WebConnConfig) *WebConn { + if cfg.Session.UserId != "" { a.Srv().Go(func() { - a.SetStatusOnline(session.UserId, false) - a.UpdateLastActivityAtIfNeeded(session) + a.SetStatusOnline(cfg.Session.UserId, false) + a.UpdateLastActivityAtIfNeeded(cfg.Session) }) } // Disable TCP_NO_DELAY for higher throughput // Unfortunately, it doesn't work for tls.Conn, // and currently, the API doesn't expose the underlying TCP conn. - tcpConn, ok := ws.UnderlyingConn().(*net.TCPConn) + tcpConn, ok := cfg.WebSocket.UnderlyingConn().(*net.TCPConn) if ok { err := tcpConn.SetNoDelay(false) if err != nil { @@ -82,25 +151,34 @@ func (a *App) NewWebConn(ws *websocket.Conn, session model.Session, t i18n.Trans } } + if cfg.activeQueue == nil { + cfg.activeQueue = make(chan model.WebSocketMessage, sendQueueSize) + } + + if cfg.deadQueue == nil && *a.srv.Config().ServiceSettings.EnableReliableWebSockets { + cfg.deadQueue = make([]*model.WebSocketEvent, deadQueueSize) + } + wc := &WebConn{ App: a, - send: make(chan model.WebSocketMessage, sendQueueSize), - WebSocket: ws, + send: cfg.activeQueue, + deadQueue: cfg.deadQueue, + deadQueuePointer: cfg.deadQueuePointer, + Sequence: int64(cfg.sequence), + WebSocket: cfg.WebSocket, lastUserActivityAt: model.GetMillis(), - UserId: session.UserId, - T: t, - Locale: locale, + UserId: cfg.Session.UserId, + T: cfg.TFunc, + Locale: cfg.Locale, + active: cfg.Active, endWritePump: make(chan struct{}), pumpFinished: make(chan struct{}), } - if *a.srv.Config().ServiceSettings.EnableReliableWebSockets { - wc.deadQueue = make([]model.WebSocketMessage, deadQueueSize) - } - - wc.SetSession(&session) - wc.SetSessionToken(session.Token) - wc.SetSessionExpiresAt(session.ExpiresAt) + wc.SetSession(&cfg.Session) + wc.SetSessionToken(cfg.Session.Token) + wc.SetSessionExpiresAt(cfg.Session.ExpiresAt) + wc.SetConnectionID(cfg.ConnectionID) return wc } @@ -136,6 +214,22 @@ func (wc *WebConn) SetConnectionID(id string) { wc.connectionID.Store(id) } +// GetConnectionID returns the connection id of the connection. +func (wc *WebConn) GetConnectionID() string { + return wc.connectionID.Load().(string) +} + +// areAllInactive returns whether all of the connections +// are inactive or not. +func areAllInactive(conns []*WebConn) bool { + for _, conn := range conns { + if conn.active { + return false + } + } + return true +} + // GetSession returns the session of the connection. func (wc *WebConn) GetSession() *model.Session { return wc.session.Load().(*model.Session) @@ -153,6 +247,8 @@ func (wc *WebConn) SetSession(v *model.Session) { // Pump starts the WebConn instance. After this, the websocket // is ready to send/receive messages. func (wc *WebConn) Pump() { + defer ReturnSessionToPool(wc.GetSession()) + var wg sync.WaitGroup wg.Add(1) go func() { @@ -164,14 +260,6 @@ func (wc *WebConn) Pump() { wg.Wait() wc.App.HubUnregister(wc) close(wc.pumpFinished) - - // TODO: - // Check if the channel is closed or not, - // if closed, then remove the entry from conn manager - // else - // take both channels, and store them in connection manager. - - defer ReturnSessionToPool(wc.GetSession()) } func (wc *WebConn) readPump() { @@ -210,6 +298,31 @@ func (wc *WebConn) writePump() { wc.WebSocket.Close() }() + if *wc.App.srv.Config().ServiceSettings.EnableReliableWebSockets && wc.Sequence != 0 { + if ok, index := wc.isInDeadQueue(wc.Sequence); ok { + if err := wc.drainDeadQueue(index); err != nil { + wc.logSocketErr("websocket.drainDeadQueue", err) + return + } + } else if wc.hasMsgLoss() { + // If the seq number is not in dead queue, but it was supposed to be, + // then generate a different connection ID, + // and set sequence to 0, and clear dead queue. + // TODO: Add metrics for this. (both true and false cases) + wc.clearDeadQueue() + wc.SetConnectionID(model.NewId()) + wc.Sequence = 0 + + // Send hello message + msg := wc.createHelloMessage() + wc.addToDeadQueue(msg) + if err := wc.writeMessage(msg); err != nil { + wc.logSocketErr("websocket.sendHello", err) + return + } + } + } + var buf bytes.Buffer // 2k is seen to be a good heuristic under which 98.5% of message sizes remain. buf.Grow(1024 * 2) @@ -219,7 +332,7 @@ func (wc *WebConn) writePump() { select { case msg, ok := <-wc.send: if !ok { - wc.writeMessage(websocket.CloseMessage, []byte{}) + wc.writeMessageBuf(websocket.CloseMessage, []byte{}) return } @@ -249,8 +362,8 @@ func (wc *WebConn) writePump() { buf.Reset() var err error if evtOk { - cpyEvt := evt.SetSequence(wc.Sequence) - err = cpyEvt.Encode(enc) + evt = evt.SetSequence(wc.Sequence) + err = evt.Encode(enc) wc.Sequence++ } else { err = enc.Encode(msg) @@ -273,11 +386,12 @@ func (wc *WebConn) writePump() { mlog.Warn("websocket.full", logData...) } - if *wc.App.srv.Config().ServiceSettings.EnableReliableWebSockets { - wc.addToDeadQueue(msg) + if *wc.App.srv.Config().ServiceSettings.EnableReliableWebSockets && + evtOk { + wc.addToDeadQueue(evt) } - if err := wc.writeMessage(websocket.TextMessage, buf.Bytes()); err != nil { + if err := wc.writeMessageBuf(websocket.TextMessage, buf.Bytes()); err != nil { wc.logSocketErr("websocket.send", err) return } @@ -286,7 +400,7 @@ func (wc *WebConn) writePump() { wc.App.Metrics().IncrementWebSocketBroadcast(msg.EventType()) } case <-ticker.C: - if err := wc.writeMessage(websocket.PingMessage, []byte{}); err != nil { + if err := wc.writeMessageBuf(websocket.PingMessage, []byte{}); err != nil { wc.logSocketErr("websocket.ticker", err) return } @@ -304,19 +418,117 @@ func (wc *WebConn) writePump() { } } -// writeMessage is a helper utility that wraps the write to the socket +// writeMessageBuf is a helper utility that wraps the write to the socket // along with setting the write deadline. -func (wc *WebConn) writeMessage(msgType int, data []byte) error { +func (wc *WebConn) writeMessageBuf(msgType int, data []byte) error { wc.WebSocket.SetWriteDeadline(time.Now().Add(writeWaitTime)) return wc.WebSocket.WriteMessage(msgType, data) } +func (wc *WebConn) writeMessage(msg *model.WebSocketEvent) error { + // We don't use the encoder from the write pump because it's unwieldy to pass encoders + // around, and this is only called during initialization of the webConn. + var buf bytes.Buffer + err := msg.Encode(json.NewEncoder(&buf)) + if err != nil { + mlog.Warn("Error in encoding websocket message", mlog.Err(err)) + return nil + } + wc.Sequence++ + + return wc.writeMessageBuf(websocket.TextMessage, buf.Bytes()) +} + // addToDeadQueue appends a message to the dead queue. -func (wc *WebConn) addToDeadQueue(msg model.WebSocketMessage) { +func (wc *WebConn) addToDeadQueue(msg *model.WebSocketEvent) { wc.deadQueue[wc.deadQueuePointer] = msg wc.deadQueuePointer = (wc.deadQueuePointer + 1) % deadQueueSize } +// hasMsgLoss indicates whether the next wanted sequence is right after +// the latest element in the dead queue, which would mean there is no message loss. +func (wc *WebConn) hasMsgLoss() bool { + var index int + if wc.deadQueuePointer == 0 { + if wc.deadQueue[deadQueueSize-1] == nil { + return false // No msg written + } + index = deadQueueSize - 1 + } else { + index = wc.deadQueuePointer - 1 + } + + if wc.deadQueue[index].GetSequence() == wc.Sequence-1 { + return false + } + return true +} + +// isInDeadQueue checks whether a given sequence number is in the dead queue or not. +// And if it is, it returns that index. +func (wc *WebConn) isInDeadQueue(seq int64) (bool, int) { + // Can be optimized to traverse backwards from deadQueuePointer + // Hopefully, traversing 128 elements is not too much overhead. + for i := 0; i < deadQueueSize; i++ { + elem := wc.deadQueue[i] + if elem == nil { + return false, 0 + } + + if elem.GetSequence() == seq { + return true, i + } + } + return false, 0 +} + +func (wc *WebConn) clearDeadQueue() { + for i := 0; i < deadQueueSize; i++ { + if wc.deadQueue[i] == nil { + return + } + wc.deadQueue[i] = nil + } + wc.deadQueuePointer = 0 +} + +// drainDeadQueue will write all messages from a given index to the socket. +// It is called with the assumption that the item with wc.Sequence is present +// in it, because otherwise it would have been cleared from WebConn. +func (wc *WebConn) drainDeadQueue(index int) error { + if wc.deadQueue[0] == nil { + // Empty queue + return nil + } + + // This means pointer hasn't rolled over. + if wc.deadQueue[wc.deadQueuePointer] == nil { + // Clear till the end of queue. + for i := index; i < wc.deadQueuePointer; i++ { + if err := wc.writeMessage(wc.deadQueue[i]); err != nil { + return err + } + } + return nil + } + + // We go on until next sequence number is smaller than previous one. + // Which means it has rolled over. + currPtr := index + for { + if err := wc.writeMessage(wc.deadQueue[currPtr]); err != nil { + return err + } + oldSeq := wc.deadQueue[currPtr].GetSequence() // TODO: possibly move this + currPtr = (currPtr + 1) % deadQueueSize // to for loop condition + newSeq := wc.deadQueue[currPtr].GetSequence() + if oldSeq > newSeq { + break + } + } + return nil +} + // InvalidateCache resets all internal data of the WebConn. func (wc *WebConn) InvalidateCache() { wc.allChannelMembers = nil diff --git a/app/web_conn_test.go b/app/web_conn_test.go index eceb5c38a9..b49889e50d 100644 --- a/app/web_conn_test.go +++ b/app/web_conn_test.go @@ -4,6 +4,10 @@ package app import ( + "bytes" + "net" + "net/http" + "net/http/httptest" "testing" "github.com/gorilla/websocket" @@ -100,17 +104,15 @@ func TestWebConnShouldSendEvent(t *testing.T) { assert.False(t, basicUserWc.shouldSendEvent(event3)) } -func TestWebConnDeadQueue(t *testing.T) { +func TestWebConnAddDeadQueue(t *testing.T) { th := Setup(t) defer th.TearDown() th.App.UpdateConfig(func(cfg *model.Config) { *cfg.ServiceSettings.EnableReliableWebSockets = true }) - session := model.Session{ - Id: model.NewId(), - } - - wc := th.App.NewWebConn(&websocket.Conn{}, session, nil, "") + wc := th.App.NewWebConn(&WebConnConfig{ + WebSocket: &websocket.Conn{}, + }) for i := 0; i < 2; i++ { msg := &model.WebSocketEvent{} @@ -119,7 +121,7 @@ func TestWebConnDeadQueue(t *testing.T) { } for i := 0; i < 2; i++ { - assert.Equal(t, int64(i), wc.deadQueue[i].(*model.WebSocketEvent).GetSequence()) + assert.Equal(t, int64(i), wc.deadQueue[i].GetSequence()) } // Should push out the first two elements @@ -129,6 +131,171 @@ func TestWebConnDeadQueue(t *testing.T) { wc.addToDeadQueue(msg) } for i := 0; i < deadQueueSize; i++ { - assert.Equal(t, int64(i+2), wc.deadQueue[(i+2)%deadQueueSize].(*model.WebSocketEvent).GetSequence()) + assert.Equal(t, int64(i+2), wc.deadQueue[(i+2)%deadQueueSize].GetSequence()) } } + +func TestWebConnIsInDeadQueue(t *testing.T) { + th := Setup(t) + defer th.TearDown() + + th.App.UpdateConfig(func(cfg *model.Config) { + *cfg.ServiceSettings.EnableReliableWebSockets = true + }) + + wc := th.App.NewWebConn(&WebConnConfig{ + WebSocket: &websocket.Conn{}, + }) + + var i int + for ; i < 2; i++ { + msg := &model.WebSocketEvent{} + msg = msg.SetSequence(int64(i)) + wc.addToDeadQueue(msg) + } + + wc.Sequence = int64(0) + ok, ind := wc.isInDeadQueue(wc.Sequence) + assert.True(t, ok) + assert.Equal(t, 0, ind) + assert.True(t, wc.hasMsgLoss()) + wc.Sequence = int64(1) + ok, ind = wc.isInDeadQueue(wc.Sequence) + assert.True(t, ok) + assert.Equal(t, 1, ind) + assert.True(t, wc.hasMsgLoss()) + wc.Sequence = int64(2) + ok, ind = wc.isInDeadQueue(wc.Sequence) + assert.False(t, ok) + assert.Equal(t, 0, ind) + assert.False(t, wc.hasMsgLoss()) + + for ; i < deadQueueSize+2; i++ { + msg := &model.WebSocketEvent{} + msg = msg.SetSequence(int64(i)) + wc.addToDeadQueue(msg) + } + + wc.Sequence = int64(129) + ok, ind = wc.isInDeadQueue(wc.Sequence) + assert.True(t, ok) + assert.Equal(t, 1, ind) + wc.Sequence = int64(128) + ok, ind = wc.isInDeadQueue(wc.Sequence) + assert.True(t, ok) + assert.Equal(t, 0, ind) + wc.Sequence = int64(2) + ok, ind = wc.isInDeadQueue(wc.Sequence) + assert.True(t, ok) + assert.Equal(t, 2, ind) + assert.True(t, wc.hasMsgLoss()) + wc.Sequence = int64(0) + ok, ind = wc.isInDeadQueue(wc.Sequence) + assert.False(t, ok) + assert.Equal(t, 0, ind) + wc.Sequence = int64(130) + ok, ind = wc.isInDeadQueue(wc.Sequence) + assert.False(t, ok) + assert.Equal(t, 0, ind) + assert.False(t, wc.hasMsgLoss()) +} + +func TestWebConnDrainDeadQueue(t *testing.T) { + th := Setup(t) + defer th.TearDown() + + th.App.UpdateConfig(func(cfg *model.Config) { + *cfg.ServiceSettings.EnableReliableWebSockets = true + }) + + var dialConn = func(t *testing.T, a *App, addr net.Addr) *WebConn { + d := websocket.Dialer{} + c, _, err := d.Dial("ws://"+addr.String()+"/ws", nil) + require.NoError(t, err) + + cfg := &WebConnConfig{ + WebSocket: c, + } + return a.NewWebConn(cfg) + } + + t.Run("Empty Queue", func(t *testing.T) { + var handler = func(t *testing.T) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + upgrader := &websocket.Upgrader{} + conn, err := upgrader.Upgrade(w, req, nil) + cnt := 0 + for err == nil { + _, _, err = conn.ReadMessage() + cnt++ + } + assert.Equal(t, 1, cnt) + if _, ok := err.(*websocket.CloseError); !ok { + require.NoError(t, err) + } + } + } + s := httptest.NewServer(handler(t)) + defer s.Close() + + wc := dialConn(t, th.App, s.Listener.Addr()) + defer wc.WebSocket.Close() + wc.clearDeadQueue() + + err := wc.drainDeadQueue(0) + require.NoError(t, err) + }) + + var handler = func(t *testing.T, seqNum int64, limit int) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + upgrader := &websocket.Upgrader{} + conn, err := upgrader.Upgrade(w, req, nil) + var buf []byte + i := seqNum + for err == nil { + _, buf, err = conn.ReadMessage() + ev := model.WebSocketEventFromJson(bytes.NewReader(buf)) + require.LessOrEqual(t, int(i), limit) + assert.Equal(t, i, ev.Sequence) + i++ + } + if _, ok := err.(*websocket.CloseError); !ok { + require.NoError(t, err) + } + } + } + + run := func(seqNum int64, limit int) { + s := httptest.NewServer(handler(t, seqNum, limit)) + defer s.Close() + + wc := dialConn(t, th.App, s.Listener.Addr()) + defer wc.WebSocket.Close() + + for i := 0; i < limit; i++ { + msg := model.NewWebSocketEvent("", "", "", "", map[string]bool{}) + msg = msg.SetSequence(int64(i)) + wc.addToDeadQueue(msg) + } + wc.Sequence = seqNum + ok, index := wc.isInDeadQueue(wc.Sequence) + require.True(t, ok) + + err := wc.drainDeadQueue(index) + require.NoError(t, err) + } + + t.Run("Half-full Queue", func(t *testing.T) { + t.Run("Middle", func(t *testing.T) { run(int64(2), 10) }) + t.Run("Beginning", func(t *testing.T) { run(int64(0), 10) }) + t.Run("End", func(t *testing.T) { run(int64(9), 10) }) + t.Run("Full", func(t *testing.T) { run(int64(deadQueueSize-1), deadQueueSize) }) + }) + + t.Run("Cycled Queue", func(t *testing.T) { + t.Run("First un-overwritten", func(t *testing.T) { run(int64(10), deadQueueSize+10) }) + t.Run("End", func(t *testing.T) { run(int64(127), deadQueueSize+10) }) + t.Run("Cycled End", func(t *testing.T) { run(int64(137), deadQueueSize+10) }) + t.Run("Overwritten First", func(t *testing.T) { run(int64(128), deadQueueSize+10) }) + }) +} diff --git a/app/web_hub.go b/app/web_hub.go index b11f16cab8..7cfc959e34 100644 --- a/app/web_hub.go +++ b/app/web_hub.go @@ -9,13 +9,15 @@ import ( "runtime/debug" "strconv" "sync/atomic" + "time" "github.com/mattermost/mattermost-server/v5/model" "github.com/mattermost/mattermost-server/v5/shared/mlog" ) const ( - broadcastQueueSize = 4096 + broadcastQueueSize = 4096 + inactiveConnReaperInterval = 5 * time.Minute ) type webConnActivityMessage struct { @@ -35,6 +37,12 @@ type webConnSessionMessage struct { isRegistered chan bool } +type webConnCheckMessage struct { + userID string + connectionID string + result chan *CheckConnResult +} + // Hub is the central place to manage all websocket connections in the server. // It handles different websocket events and sending messages to individual // user connections. @@ -54,6 +62,7 @@ type Hub struct { directMsg chan *webConnDirectMessage explicitStop bool checkRegistered chan *webConnSessionMessage + checkConn chan *webConnCheckMessage } // NewWebHub creates a new Hub. @@ -69,6 +78,7 @@ func (a *App) NewWebHub() *Hub { activity: make(chan *webConnActivityMessage), directMsg: make(chan *webConnDirectMessage), checkRegistered: make(chan *webConnSessionMessage), + checkConn: make(chan *webConnCheckMessage), } } @@ -289,6 +299,14 @@ func (a *App) SessionIsRegistered(session model.Session) bool { return false } +func (a *App) CheckWebConn(userID, connectionID string) *CheckConnResult { + hub := a.GetHubForUserId(userID) + if hub != nil { + return hub.CheckConn(userID, connectionID) + } + return nil +} + // Register registers a connection to the hub. func (h *Hub) Register(webConn *WebConn) { select { @@ -320,6 +338,20 @@ func (h *Hub) IsRegistered(userID, sessionToken string) bool { return false } +func (h *Hub) CheckConn(userID, connectionID string) *CheckConnResult { + req := &webConnCheckMessage{ + userID: userID, + connectionID: connectionID, + result: make(chan *CheckConnResult), + } + select { + case h.checkConn <- req: + return <-req.result + case <-h.stop: + } + return nil +} + // Broadcast broadcasts the message to all connections in the hub. func (h *Hub) Broadcast(message *model.WebSocketEvent) { // XXX: The hub nil check is because of the way we setup our tests. We call @@ -387,7 +419,10 @@ func (h *Hub) Start() { doStart = func() { mlog.Debug("Hub is starting", mlog.Int("index", h.connectionIndex)) - connIndex := newHubConnectionIndex() + ticker := time.NewTicker(inactiveConnReaperInterval) + defer ticker.Stop() + + connIndex := newHubConnectionIndex(inactiveConnReaperInterval) for { select { @@ -395,27 +430,70 @@ func (h *Hub) Start() { conns := connIndex.ForUser(webSessionMessage.userID) var isRegistered bool for _, conn := range conns { + if !conn.active { + continue + } if conn.GetSessionToken() == webSessionMessage.sessionToken { isRegistered = true } } webSessionMessage.isRegistered <- isRegistered + case req := <-h.checkConn: + var res *CheckConnResult + conn := connIndex.GetInactiveByConnectionID(req.userID, req.connectionID) + if conn != nil { + res = &CheckConnResult{ + ConnectionID: req.connectionID, + UserID: req.userID, + ActiveQueue: conn.send, + DeadQueue: conn.deadQueue, + DeadQueuePointer: conn.deadQueuePointer, + } + } + req.result <- res + case <-ticker.C: + connIndex.RemoveInactiveConnections() case webConn := <-h.register: + var oldConn *WebConn + if *h.app.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. + webConn.active = true + connIndex.Add(webConn) - atomic.StoreInt64(&h.connectionCount, int64(len(connIndex.All()))) - if webConn.IsAuthenticated() { + 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. + // 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. webConn.send <- webConn.createHelloMessage() } case webConn := <-h.unregister: - connIndex.Remove(webConn) - atomic.StoreInt64(&h.connectionCount, int64(len(connIndex.All()))) + // If already removed (via queue full), then removing again becomes a noop. + // But if not removed, mark inactive. + if *h.app.Config().ServiceSettings.EnableReliableWebSockets { + webConn.active = false + } else { + connIndex.Remove(webConn) + } + + atomic.StoreInt64(&h.connectionCount, int64(connIndex.AllActive())) if webConn.UserId == "" { continue } conns := connIndex.ForUser(webConn.UserId) - if len(conns) == 0 { + if len(conns) == 0 || areAllInactive(conns) { h.app.Srv().Go(func() { h.app.SetStatusOffline(webConn.UserId, false) }) @@ -423,6 +501,9 @@ func (h *Hub) Start() { } var latestActivity int64 = 0 for _, conn := range conns { + if !conn.active { + continue + } if conn.lastUserActivityAt > latestActivity { latestActivity = conn.lastUserActivityAt } @@ -439,6 +520,9 @@ func (h *Hub) Start() { } case activity := <-h.activity: for _, webConn := range connIndex.ForUser(activity.userID) { + if !webConn.active { + continue + } if webConn.GetSessionToken() == activity.sessionToken { webConn.lastUserActivityAt = activity.activityAt } @@ -531,12 +615,16 @@ type hubConnectionIndex struct { // byConnection serves the dual purpose of storing the index of the webconn // in the value of byUserId map, and also to get all connections. byConnection map[*WebConn]int + // staleThreshold is the limit beyond which inactive connections + // will be deleted. + staleThreshold time.Duration } -func newHubConnectionIndex() *hubConnectionIndex { +func newHubConnectionIndex(interval time.Duration) *hubConnectionIndex { return &hubConnectionIndex{ - byUserId: make(map[string][]*WebConn), - byConnection: make(map[*WebConn]int), + byUserId: make(map[string][]*WebConn), + byConnection: make(map[*WebConn]int), + staleThreshold: interval, } } @@ -570,10 +658,67 @@ func (i *hubConnectionIndex) Has(wc *WebConn) bool { return ok } +// ForUser returns all connections for a user ID. func (i *hubConnectionIndex) ForUser(id string) []*WebConn { return i.byUserId[id] } +// All returns the full webConn index. 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 { + // To handle empty sessions. + if userID == "" { + return nil + } + for _, conn := range i.ForUser(userID) { + if conn.GetConnectionID() == connectionID && !conn.active { + i.Remove(conn) + return conn + } + } + return nil +} + +// RemoveInactiveConnections removes all inactive connections whose lastUserActivityAt +// exceeded staleThreshold. +func (i *hubConnectionIndex) RemoveInactiveConnections() { + now := model.GetMillis() + for conn := range i.byConnection { + if !conn.active && now-conn.lastUserActivityAt > i.staleThreshold.Milliseconds() { + i.Remove(conn) + } + } +} + +// AllActive returns the number of active connections. +// This is only called during register/unregister so we can take +// a bit of perf hit here. +func (i *hubConnectionIndex) AllActive() int { + cnt := 0 + for conn := range i.byConnection { + if conn.active { + cnt++ + } + } + return cnt +} diff --git a/app/web_hub_test.go b/app/web_hub_test.go index c979dd182b..d95705d99b 100644 --- a/app/web_hub_test.go +++ b/app/web_hub_test.go @@ -46,7 +46,13 @@ func registerDummyWebConn(t *testing.T, a *App, addr net.Addr, userID string) *W c, _, err := d.Dial("ws://"+addr.String()+"/ws", nil) require.NoError(t, err) - wc := a.NewWebConn(c, *session, i18n.IdentityTfunc(), "en") + cfg := &WebConnConfig{ + WebSocket: c, + Session: *session, + TFunc: i18n.IdentityTfunc(), + Locale: "en", + } + wc := a.NewWebConn(cfg) a.HubRegister(wc) go wc.Pump() return wc @@ -199,7 +205,7 @@ func TestHubConnIndex(t *testing.T) { th := Setup(t) defer th.TearDown() - connIndex := newHubConnectionIndex() + connIndex := newHubConnectionIndex(1 * time.Second) // User1 wc1 := &WebConn{ @@ -270,6 +276,56 @@ func TestHubConnIndex(t *testing.T) { }) } +func TestHubConnIndexInactive(t *testing.T) { + connIndex := newHubConnectionIndex(2 * time.Second) + + // User1 + wc1 := &WebConn{ + UserId: model.NewId(), + active: true, + } + wc1.SetConnectionID("conn1") + + // User2 + wc2 := &WebConn{ + UserId: model.NewId(), + active: true, + } + wc2.SetConnectionID("conn2") + wc3 := &WebConn{ + UserId: wc2.UserId, + active: false, + } + wc3.SetConnectionID("conn3") + + connIndex.Add(wc1) + 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")) + assert.False(t, connIndex.Has(wc3)) + assert.Len(t, connIndex.ForUser(wc2.UserId), 1) + + wc3.lastUserActivityAt = model.GetMillis() + connIndex.Add(wc3) + connIndex.RemoveInactiveConnections() + assert.True(t, connIndex.Has(wc3)) + assert.Len(t, connIndex.ForUser(wc2.UserId), 2) + assert.Len(t, connIndex.All(), 3) + + wc3.lastUserActivityAt = model.GetMillis() - (time.Minute).Milliseconds() + connIndex.RemoveInactiveConnections() + assert.False(t, connIndex.Has(wc3)) + assert.Len(t, connIndex.ForUser(wc2.UserId), 1) + assert.Len(t, connIndex.All(), 2) +} + func TestHubIsRegistered(t *testing.T) { th := Setup(t).InitBasic() defer th.TearDown() @@ -303,7 +359,7 @@ func TestHubIsRegistered(t *testing.T) { func BenchmarkHubConnIndex(b *testing.B) { th := Setup(b).InitBasic() defer th.TearDown() - connIndex := newHubConnectionIndex() + connIndex := newHubConnectionIndex(1 * time.Second) // User1 wc1 := &WebConn{ diff --git a/config/client.go b/config/client.go index 5d0acb7d9c..9154c4f266 100644 --- a/config/client.go +++ b/config/client.go @@ -101,6 +101,8 @@ func GenerateClientConfig(c *model.Config, telemetryID string, license *model.Li props["EnableLegacySidebar"] = strconv.FormatBool(*c.ServiceSettings.EnableLegacySidebar) + props["EnableReliableWebSockets"] = strconv.FormatBool(*c.ServiceSettings.EnableReliableWebSockets) + // Set default values for all options that require a license. props["ExperimentalHideTownSquareinLHS"] = "false" props["ExperimentalTownSquareIsReadOnly"] = "false"