From 5215be51be5105baf1591bf4d92f553e13f1a1d1 Mon Sep 17 00:00:00 2001
From: Agniva De Sarker <agnivade@yahoo.co.in>
Date: Thu, 1 Apr 2021 15:31:19 +0530
Subject: [PATCH] MM-34487: Add dead queue (#17307)

* MM-34487: Add dead queue

Just a circular buffer to store dead messages for now.
Not controlling this via config flag because this does
not have any effect except taking some more memory
per connection

```release-note
NONE
```

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

* Wrap with config

* fix test
---
 app/web_conn.go      | 36 ++++++++++++++++++++++++++++--------
 app/web_conn_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/app/web_conn.go b/app/web_conn.go
index 60f03124f2..c78f53c2ec 100644
--- a/app/web_conn.go
+++ b/app/web_conn.go
@@ -29,10 +29,11 @@ const (
 	pingInterval           = (pongWaitTime * 6) / 10
 	authCheckInterval      = 5 * time.Second
 	webConnMemberCacheTime = 1000 * 60 * 30 // 30 minutes
+	deadQueueSize          = 128            // Approximated from /proc/sys/net/core/wmem_default / 2048 (avg msg size)
 )
 
 // WebConn represents a single websocket connection to a user.
-// It contains all the necesarry state to manage sending/receiving data to/from
+// It contains all the necessary state to manage sending/receiving data to/from
 // a websocket.
 type WebConn struct {
 	sessionExpiresAt int64 // This should stay at the top for 64-bit alignment of 64-bit words accessed atomically
@@ -47,11 +48,18 @@ type WebConn struct {
 	lastAllChannelMembersTime int64
 	lastUserActivityAt        int64
 	send                      chan model.WebSocketMessage
-	sessionToken              atomic.Value
-	session                   atomic.Value
-	connectionID              atomic.Value
-	endWritePump              chan struct{}
-	pumpFinished              chan struct{}
+	// deadQueue behaves like a queue of a finite size
+	// which is used to store all messages that are sent via the websocket.
+	// 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{}
 }
 
 // NewWebConn returns a new WebConn instance.
@@ -86,6 +94,10 @@ func (a *App) NewWebConn(ws *websocket.Conn, session model.Session, t i18n.Trans
 		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)
@@ -261,13 +273,15 @@ func (wc *WebConn) writePump() {
 				mlog.Warn("websocket.full", logData...)
 			}
 
+			if *wc.App.srv.Config().ServiceSettings.EnableReliableWebSockets {
+				wc.addToDeadQueue(msg)
+			}
+
 			if err := wc.writeMessage(websocket.TextMessage, buf.Bytes()); err != nil {
 				wc.logSocketErr("websocket.send", err)
 				return
 			}
 
-			// TODO: move to dead queue
-
 			if wc.App.Metrics() != nil {
 				wc.App.Metrics().IncrementWebSocketBroadcast(msg.EventType())
 			}
@@ -297,6 +311,12 @@ func (wc *WebConn) writeMessage(msgType int, data []byte) error {
 	return wc.WebSocket.WriteMessage(msgType, data)
 }
 
+// addToDeadQueue appends a message to the dead queue.
+func (wc *WebConn) addToDeadQueue(msg model.WebSocketMessage) {
+	wc.deadQueue[wc.deadQueuePointer] = msg
+	wc.deadQueuePointer = (wc.deadQueuePointer + 1) % deadQueueSize
+}
+
 // 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 d289ee9f1c..eceb5c38a9 100644
--- a/app/web_conn_test.go
+++ b/app/web_conn_test.go
@@ -6,6 +6,7 @@ package app
 import (
 	"testing"
 
+	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -98,3 +99,36 @@ func TestWebConnShouldSendEvent(t *testing.T) {
 	event3 := model.NewWebSocketEvent(model.WEBSOCKET_EVENT_UPDATE_TEAM, "wrongId", "", "", nil)
 	assert.False(t, basicUserWc.shouldSendEvent(event3))
 }
+
+func TestWebConnDeadQueue(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, "")
+
+	for i := 0; i < 2; i++ {
+		msg := &model.WebSocketEvent{}
+		msg = msg.SetSequence(int64(i))
+		wc.addToDeadQueue(msg)
+	}
+
+	for i := 0; i < 2; i++ {
+		assert.Equal(t, int64(i), wc.deadQueue[i].(*model.WebSocketEvent).GetSequence())
+	}
+
+	// Should push out the first two elements
+	for i := 0; i < deadQueueSize; i++ {
+		msg := &model.WebSocketEvent{}
+		msg = msg.SetSequence(int64(i + 2))
+		wc.addToDeadQueue(msg)
+	}
+	for i := 0; i < deadQueueSize; i++ {
+		assert.Equal(t, int64(i+2), wc.deadQueue[(i+2)%deadQueueSize].(*model.WebSocketEvent).GetSequence())
+	}
+}