MM-57152: Get webconn count from the whole cluster (#26813)

We were setting the user status to offline without
checking for connections on other nodes in a cluster.

Now we implement a request-response mechanism for the whole
cluster and we check that before setting a user to offline.

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

```release-note
Fix a bug where the user status would incorrectly be set to offline
without checking for connections in other nodes in an HA cluster.
```

Co-authored-by: Ibrahim Serdar Acikgoz <serdaracikgoz86@gmail.com>
Co-authored-by: Mattermost Build <build@mattermost.com>
Co-authored-by: Ben Schumacher <ben.schumacher@mattermost.com>

server/channels/app/busy_test.go

@@ -151,3 +151,6 @@ func (c *ClusterMock) ConfigChanged(previousConfig *model.Config, newConfig *mod
 	return nil
 }
 func (c *ClusterMock) HealthScore() int { return 0 }
+func (c *ClusterMock) WebConnCountForUser(userID string) (int, *model.AppError) {
+	return 0, nil
+}

server/channels/app/platform/busy_test.go

@@ -149,3 +149,6 @@ func (c *ClusterMock) ConfigChanged(previousConfig *model.Config, newConfig *mod
 	return nil
 }
 func (c *ClusterMock) HealthScore() int { return 0 }
+func (c *ClusterMock) WebConnCountForUser(userID string) (int, *model.AppError) {
+	return 0, nil
+}

server/channels/app/platform/web_conn.go

@@ -372,17 +372,6 @@ func (wc *WebConn) isSet(val string) bool {
 	return val != UnsetPresenceIndicator
 }
 
-// areAllInactive returns whether all of the connections
-// are inactive or not.
-func areAllInactive(conns []*WebConn) bool {
-	for _, conn := range conns {
-		if conn.active.Load() {
-			return false
-		}
-	}
-	return true
-}
-
 // GetSession returns the session of the connection.
 func (wc *WebConn) GetSession() *model.Session {
 	return wc.session.Load()

server/channels/app/platform/web_hub.go

@@ -50,6 +50,11 @@ type webConnCheckMessage struct {
 	result       chan *CheckConnResult
 }
 
+type webConnCountMessage struct {
+	userID string
+	result chan int
+}
+
 // 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.
@@ -70,6 +75,7 @@ type Hub struct {
 	explicitStop    bool
 	checkRegistered chan *webConnSessionMessage
 	checkConn       chan *webConnCheckMessage
+	connCount       chan *webConnCountMessage
 	broadcastHooks  map[string]BroadcastHook
 }
 
@@ -87,6 +93,7 @@ func newWebHub(ps *PlatformService) *Hub {
 		directMsg:       make(chan *webConnDirectMessage),
 		checkRegistered: make(chan *webConnSessionMessage),
 		checkConn:       make(chan *webConnCheckMessage),
+		connCount:       make(chan *webConnCountMessage),
 	}
 }
 
@@ -236,6 +243,16 @@ func (ps *PlatformService) CheckWebConn(userID, connectionID string) *CheckConnR
 	return nil
 }
 
+// WebConnCountForUser returns the number of active websocket connections
+// for a given userID.
+func (ps *PlatformService) WebConnCountForUser(userID string) int {
+	hub := ps.GetHubForUserId(userID)
+	if hub != nil {
+		return hub.WebConnCountForUser(userID)
+	}
+	return 0
+}
+
 // Register registers a connection to the hub.
 func (h *Hub) Register(webConn *WebConn) {
 	select {
@@ -281,6 +298,19 @@ func (h *Hub) CheckConn(userID, connectionID string) *CheckConnResult {
 	return nil
 }
 
+func (h *Hub) WebConnCountForUser(userID string) int {
+	req := &webConnCountMessage{
+		userID: userID,
+		result: make(chan int),
+	}
+	select {
+	case h.connCount <- req:
+		return <-req.result
+	case <-h.stop:
+	}
+	return 0
+}
+
 // 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
@@ -381,6 +411,8 @@ func (h *Hub) Start() {
 					}
 				}
 				req.result <- res
+			case req := <-h.connCount:
+				req.result <- connIndex.ForUserActiveCount(req.userID)
 			case <-ticker.C:
 				connIndex.RemoveInactiveConnections()
 			case webConn := <-h.register:
@@ -420,7 +452,26 @@ func (h *Hub) Start() {
 				if len(conns) == 0 || areAllInactive(conns) {
 					userID := webConn.UserId
 					h.platform.Go(func() {
-						h.platform.SetStatusOffline(userID, false)
+						// If this is an HA setup, get count for this user
+						// from other nodes.
+						var clusterCnt int
+						var appErr *model.AppError
+						if h.platform.Cluster() != nil {
+							clusterCnt, appErr = h.platform.Cluster().WebConnCountForUser(userID)
+						}
+						if appErr != nil {
+							mlog.Error("Error in trying to get the webconn count from cluster", mlog.Err(appErr))
+							// We take a conservative approach
+							// and do not set status to offline in case
+							// there's an error, rather than potentially
+							// incorrectly setting status to offline.
+							return
+						}
+						// Only set to offline if there are no
+						// active connections in other nodes as well.
+						if clusterCnt == 0 {
+							h.platform.SetStatusOffline(userID, false)
+						}
 					})
 					continue
 				}
@@ -552,6 +603,17 @@ func (h *Hub) Start() {
 	go doRecoverableStart()
 }
 
+// areAllInactive returns whether all of the connections
+// are inactive or not.
+func areAllInactive(conns []*WebConn) bool {
+	for _, conn := range conns {
+		if conn.active.Load() {
+			return false
+		}
+	}
+	return true
+}
+
 // hubConnectionIndex provides fast addition, removal, and iteration of web connections.
 // It requires 3 functionalities which need to be very fast:
 // - check if a connection exists or not.
@@ -629,6 +691,17 @@ func (i *hubConnectionIndex) ForUser(id string) []*WebConn {
 	return conns
 }
 
+// ForUserActiveCount returns the number of active connections for a userID
+func (i *hubConnectionIndex) ForUserActiveCount(id string) int {
+	cnt := 0
+	for _, conn := range i.ForUser(id) {
+		if conn.active.Load() {
+			cnt++
+		}
+	}
+	return cnt
+}
+
 // ForConnection returns the connection from its ID.
 func (i *hubConnectionIndex) ForConnection(id string) *WebConn {
 	return i.byConnectionId[id]

server/channels/app/platform/web_hub_test.go

@@ -291,6 +291,7 @@ func TestHubConnIndex(t *testing.T) {
 
 		assert.ElementsMatch(t, connIndex.ForUser(wc2.UserId), []*WebConn{wc2, wc4})
 		assert.ElementsMatch(t, connIndex.ForUser(wc1.UserId), []*WebConn{})
+		assert.Len(t, connIndex.ForUser(wc1.UserId), 0)
 		assert.Len(t, connIndex.All(), 2)
 		assert.False(t, connIndex.Has(wc1))
 		assert.True(t, connIndex.Has(wc2))
@@ -454,7 +455,9 @@ func TestHubConnIndexInactive(t *testing.T) {
 	connIndex.Add(wc3)
 
 	assert.Nil(t, connIndex.RemoveInactiveByConnectionID(wc2.UserId, "conn2"))
+	assert.Equal(t, connIndex.ForUserActiveCount(wc2.UserId), 1)
 	assert.NotNil(t, connIndex.RemoveInactiveByConnectionID(wc2.UserId, "conn3"))
+	assert.Equal(t, connIndex.ForUserActiveCount(wc2.UserId), 1)
 	assert.Nil(t, connIndex.RemoveInactiveByConnectionID(wc1.UserId, "conn3"))
 	assert.False(t, connIndex.Has(wc3))
 	assert.Len(t, connIndex.ForUser(wc2.UserId), 1)
@@ -464,12 +467,14 @@ func TestHubConnIndexInactive(t *testing.T) {
 	connIndex.RemoveInactiveConnections()
 	assert.True(t, connIndex.Has(wc3))
 	assert.Len(t, connIndex.ForUser(wc2.UserId), 2)
+	assert.Equal(t, connIndex.ForUserActiveCount(wc2.UserId), 1)
 	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.Equal(t, connIndex.ForUserActiveCount(wc2.UserId), 1)
 	assert.Len(t, connIndex.All(), 2)
 }
 
@@ -543,6 +548,35 @@ func TestHubIsRegistered(t *testing.T) {
 	assert.False(t, th.Service.SessionIsRegistered(*session4))
 }
 
+func TestHubWebConnCount(t *testing.T) {
+	th := Setup(t).InitBasic()
+	defer th.TearDown()
+
+	session, err := th.Service.CreateSession(th.Context, &model.Session{
+		UserId: th.BasicUser.Id,
+	})
+	require.NoError(t, err)
+
+	mockSuite := &platform_mocks.SuiteIFace{}
+	mockSuite.On("GetSession", session.Token).Return(session, nil)
+	th.Suite = mockSuite
+
+	s := httptest.NewServer(dummyWebsocketHandler(t))
+	defer s.Close()
+
+	th.Service.Start(nil)
+	wc1 := registerDummyWebConn(t, th, s.Listener.Addr(), session)
+	wc2 := registerDummyWebConn(t, th, s.Listener.Addr(), session)
+	defer wc1.Close()
+
+	assert.Equal(t, 2, th.Service.WebConnCountForUser(th.BasicUser.Id))
+
+	wc2.Close()
+
+	assert.Equal(t, 1, th.Service.WebConnCountForUser(th.BasicUser.Id))
+	assert.Equal(t, 0, th.Service.WebConnCountForUser("none"))
+}
+
 // Always run this with -benchtime=0.1s
 // See: https://github.com/golang/go/issues/27217.
 func BenchmarkHubConnIndex(b *testing.B) {

server/channels/testlib/cluster.go

@@ -103,3 +103,7 @@ func (c *FakeClusterInterface) ClearMessages() {
 	defer c.mut.Unlock()
 	c.messages = nil
 }
+
+func (c *FakeClusterInterface) WebConnCountForUser(userID string) (int, *model.AppError) {
+	return 0, nil
+}

server/einterfaces/cluster.go

@@ -29,4 +29,7 @@ type ClusterInterface interface {
 	QueryLogs(page, perPage int) (map[string][]string, *model.AppError)
 	GetPluginStatuses() (model.PluginStatuses, *model.AppError)
 	ConfigChanged(previousConfig *model.Config, newConfig *model.Config, sendToOtherServer bool) *model.AppError
+	// WebConnCountForUser returns the number of active webconn connections
+	// for a given userID.
+	WebConnCountForUser(userID string) (int, *model.AppError)
 }

server/einterfaces/mocks/ClusterInterface.go

@@ -301,6 +301,36 @@ func (_m *ClusterInterface) StopInterNodeCommunication() {
 	_m.Called()
 }
 
+// WebConnCountForUser provides a mock function with given fields: userID
+func (_m *ClusterInterface) WebConnCountForUser(userID string) (int, *model.AppError) {
+	ret := _m.Called(userID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for WebConnCountForUser")
+	}
+
+	var r0 int
+	var r1 *model.AppError
+	if rf, ok := ret.Get(0).(func(string) (int, *model.AppError)); ok {
+		return rf(userID)
+	}
+	if rf, ok := ret.Get(0).(func(string) int); ok {
+		r0 = rf(userID)
+	} else {
+		r0 = ret.Get(0).(int)
+	}
+
+	if rf, ok := ret.Get(1).(func(string) *model.AppError); ok {
+		r1 = rf(userID)
+	} else {
+		if ret.Get(1) != nil {
+			r1 = ret.Get(1).(*model.AppError)
+		}
+	}
+
+	return r0, r1
+}
+
 // NewClusterInterface creates a new instance of ClusterInterface. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
 // The first argument is typically a *testing.T value.
 func NewClusterInterface(t interface {

server/public/model/cluster_message.go

@@ -51,6 +51,8 @@ const (
 	ClusterGossipEventResponseGetPluginStatuses = "gossip_response_plugin_statuses"
 	ClusterGossipEventRequestSaveConfig         = "gossip_request_save_config"
 	ClusterGossipEventResponseSaveConfig        = "gossip_response_save_config"
+	ClusterGossipEventRequestWebConnCount       = "gossip_request_webconn_count"
+	ClusterGossipEventResponseWebConnCount      = "gossip_response_webconn_count"
 
 	// SendTypes for ClusterMessage.
 	ClusterSendBestEffort = "best_effort"