MM-32950: Reliable WebSockets: Basic single server (#17406)

* MM-32950: Reliable WebSockets: Basic single server

This PR adds reliable websocket support for a single server.

Below is a brief overview of the three states of a connection:

Normal:
- All messages are routed via web hub.
- Each web conn has a send queue to which it gets pushed.
- A message gets pulled from the queue, and before it
gets written to the wire, it is added to the dead queue.

Disconnect:
- Hub Unregister gets called, where the connection is just
marked as inactive. And new messages keep getting pushed
to the send queue.

If it gets full, the channel is closed and the conn gets removed
from conn index.

Reconnect:
- We query the hub for the connection ID, and get back the
queues.
- We construct a WebConn reusing the old queues, or a fresh one
depending on whether the connection ID was found or not.
- Now there is a tricky bit here which needs to be carefully processed.
On register, we would always send the hello message in the send queue.
But we cannot do that now because the send queue might already have messages.

Therefore, we don't send the hello message from web hub, if we reuse a connection.

Instead, we move that logic to the web conn write pump. We check if
the sequence number is in dead queue, and if it is, then we drain
the dead queue, and start consuming from the active queue.
No hello message is sent here.

But if the message does not exist in the dead queue, and the sequence number
is actually something that should have existed, then we set
a new connction id and clear the dead queue, and send a hello message.
The client, on receiving a new connection id will automatically
set its sequence number to 0, and make the sync API calls to manage
any lost data.

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

```release-note
NONE
```

* gofmt

* Add EnableReliableWebSockets to the client config

* Refactoring isInDeadQueue

* Passing index to drainDeadQueue

* refactoring webconn

* fix pointer

* review comments

* simplify hasMsgLoss

* safety comment

* fix test

* Trigger CI

* Trigger CI

Co-authored-by: Devin Binnie <devin.binnie@mattermost.com>
Co-authored-by: Mattermod <mattermod@users.noreply.github.com>

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)
 	}

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()

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")

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

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) })
+	})
+}

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
+}

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{

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"