MM-63130: Move to webHub iteration to be alloc-free (#30792)

We switch to using iterators introduced in Go 1.23 to make iteration alloc-free and fast. And since element removal is allowed while iterating a map, this also means we don't need to even copy the slice any more. While here, we also address the comment https://github.com/mattermost/mattermost/pull/30178#discussion_r1954862151. I have simply gone back to using []string as the map entry rather than a type alias or a redirection with a struct. https://mattermost.atlassian.net/browse/MM-63130 ```release-note NONE ``` * Changed back nil to len ```release-note NONE ``` * fixing unused assignment ```release-note NONE ``` * add benchmark ``` goos: linux goarch: amd64 pkg: github.com/mattermost/mattermost/server/v8/channels/app/platform cpu: Intel(R) Core(TM) i5-8265U CPU @ 1.60GHz │ old.txt │ new.txt │ │ sec/op │ sec/op vs base │ HubConnIndexIterator/2_users-8 93.53n ± 1% 38.09n ± 1% -59.27% (p=0.000 n=10) HubConnIndexIterator/3_users-8 106.30n ± 0% 38.41n ± 1% -63.86% (p=0.000 n=10) HubConnIndexIterator/4_users-8 111.30n ± 1% 38.66n ± 1% -65.27% (p=0.000 n=10) geomean 103.4n 38.39n -62.89% │ old.txt │ new.txt │ │ B/op │ B/op vs base │ HubConnIndexIterator/2_users-8 16.00 ± 0% 24.00 ± 0% +50.00% (p=0.000 n=10) HubConnIndexIterator/3_users-8 24.00 ± 0% 24.00 ± 0% ~ (p=1.000 n=10) ¹ HubConnIndexIterator/4_users-8 32.00 ± 0% 24.00 ± 0% -25.00% (p=0.000 n=10) geomean 23.08 24.00 +4.00% ¹ all samples are equal │ old.txt │ new.txt │ │ allocs/op │ allocs/op vs base │ HubConnIndexIterator/2_users-8 1.000 ± 0% 1.000 ± 0% ~ (p=1.000 n=10) ¹ HubConnIndexIterator/3_users-8 1.000 ± 0% 1.000 ± 0% ~ (p=1.000 n=10) ¹ HubConnIndexIterator/4_users-8 1.000 ± 0% 1.000 ± 0% ~ (p=1.000 n=10) ¹ geomean 1.000 1.000 +0.00% ¹ all samples are equal ``` ```release-note NONE ``` * ForChannel test as well ```release-note NONE ``` * review comments ```release-note NONE ``` * fix lint errors ```release-note NONE ``` --------- Co-authored-by: Mattermost Build <build@mattermost.com>
2025-05-11 12:00:12 +05:30
--- a/server/channels/app/platform/web_hub.go
+++ b/server/channels/app/platform/web_hub.go
@@ -6,6 +6,8 @@ package platform
 import (
 	"fmt"
 	"hash/maphash"
+	"iter"
+	"maps"
 	"runtime"
 	"runtime/debug"
 	"strconv"
@@ -494,9 +496,8 @@ func (h *Hub) Start() {
 		for {
 			select {
 			case webSessionMessage := <-h.checkRegistered:
-				conns := connIndex.ForUser(webSessionMessage.userID)
 				var isRegistered bool
-				for _, conn := range conns {
+				for conn := range connIndex.ForUser(webSessionMessage.userID) {
 					if !conn.Active.Load() {
 						continue
 					}
@@ -556,7 +557,9 @@ func (h *Hub) Start() {
 				}

 				conns := connIndex.ForUser(webConn.UserId)
-				if len(conns) == 0 || areAllInactive(conns) {
+				// areAllInactive also returns true if there are no connections,
+				// which is intentional.
+				if areAllInactive(conns) {
 					userID := webConn.UserId
 					h.platform.Go(func() {
 						// If this is an HA setup, get count for this user
@@ -583,7 +586,7 @@ func (h *Hub) Start() {
 					continue
 				}
 				var latestActivity int64
-				for _, conn := range conns {
+				for conn := range conns {
 					if !conn.Active.Load() {
 						continue
 					}
@@ -599,7 +602,7 @@ func (h *Hub) Start() {
 					})
 				}
 			case userID := <-h.invalidateUser:
-				for _, webConn := range connIndex.ForUser(userID) {
+				for webConn := range connIndex.ForUser(userID) {
 					webConn.InvalidateCache()
 				}

@@ -610,12 +613,12 @@ func (h *Hub) Start() {
 				err := connIndex.InvalidateCMCacheForUser(userID)
 				if err != nil {
 					h.platform.Log().Error("Error while invalidating channel member cache", mlog.String("user_id", userID), mlog.Err(err))
-					for _, webConn := range connIndex.ForUser(userID) {
+					for webConn := range connIndex.ForUser(userID) {
 						closeAndRemoveConn(connIndex, webConn)
 					}
 				}
 			case activity := <-h.activity:
-				for _, webConn := range connIndex.ForUser(activity.userID) {
+				for webConn := range connIndex.ForUser(activity.userID) {
 					if !webConn.Active.Load() {
 						continue
 					}
@@ -667,18 +670,21 @@ func (h *Hub) Start() {
 					}
 				}

-				var targetConns []*WebConn
-				if connID := msg.GetBroadcast().ConnectionId; connID != "" {
-					if webConn := connIndex.ForConnection(connID); webConn != nil {
-						targetConns = append(targetConns, webConn)
-					}
-				} else if userID := msg.GetBroadcast().UserId; userID != "" {
+				// Quick return for a single connection.
+				if webConn := connIndex.ForConnection(msg.GetBroadcast().ConnectionId); webConn != nil {
+					broadcast(webConn)
+					continue
+				}
+
+				fastIteration := *h.platform.Config().ServiceSettings.EnableWebHubChannelIteration
+				var targetConns iter.Seq[*WebConn]
+				if userID := msg.GetBroadcast().UserId; userID != "" {
 					targetConns = connIndex.ForUser(userID)
-				} else if channelID := msg.GetBroadcast().ChannelId; channelID != "" && *h.platform.Config().ServiceSettings.EnableWebHubChannelIteration {
+				} else if channelID := msg.GetBroadcast().ChannelId; channelID != "" && fastIteration {
 					targetConns = connIndex.ForChannel(channelID)
 				}
 				if targetConns != nil {
-					for _, webConn := range targetConns {
+					for webConn := range targetConns {
 						broadcast(webConn)
 					}
 					continue
@@ -688,7 +694,7 @@ func (h *Hub) Start() {
 				// method, there would be events scoped to a channel being sent to multiple hubs. And only one hub would
 				// have the targetConns. Therefore, we need to stop here if channel based iteration is enabled, and it's a
 				// channel-scoped event.
-				if channelID := msg.GetBroadcast().ChannelId; channelID != "" && *h.platform.Config().ServiceSettings.EnableWebHubChannelIteration {
+				if channelID := msg.GetBroadcast().ChannelId; channelID != "" && fastIteration {
 					continue
 				}

@@ -732,9 +738,10 @@ func (h *Hub) Start() {
 }

 // areAllInactive returns whether all of the connections
-// are inactive or not.
-func areAllInactive(conns []*WebConn) bool {
-	for _, conn := range conns {
+// are inactive or not. It also returns true if there are
+// no connections which is also intentional.
+func areAllInactive(conns iter.Seq[*WebConn]) bool {
+	for conn := range conns {
 		if conn.Active.Load() {
 			return false
 		}
@@ -749,10 +756,6 @@ func closeAndRemoveConn(connIndex *hubConnectionIndex, conn *WebConn) {
 	connIndex.Remove(conn)
 }

-type connMetadata struct {
-	channelIDs []string
-}
-
 // hubConnectionIndex provides fast addition, removal, and iteration of web connections.
 // It requires 4 functionalities which need to be very fast:
 // - check if a connection exists or not.
@@ -766,7 +769,7 @@ type hubConnectionIndex struct {
 	byChannelID map[string]map[*WebConn]struct{}
 	// byConnection serves the dual purpose of storing the channelIDs
 	// and also to get all connections
-	byConnection   map[*WebConn]connMetadata
+	byConnection   map[*WebConn][]string
 	byConnectionId map[string]*WebConn
 	// staleThreshold is the limit beyond which inactive connections
 	// will be deleted.
@@ -785,7 +788,7 @@ func newHubConnectionIndex(interval time.Duration,
 	return &hubConnectionIndex{
 		byUserId:       make(map[string]map[*WebConn]struct{}),
 		byChannelID:    make(map[string]map[*WebConn]struct{}),
-		byConnection:   make(map[*WebConn]connMetadata),
+		byConnection:   make(map[*WebConn][]string),
 		byConnectionId: make(map[string]*WebConn),
 		staleThreshold: interval,
 		store:          store,
@@ -820,15 +823,13 @@ func (i *hubConnectionIndex) Add(wc *WebConn) error {
 		i.byUserId[wc.UserId] = make(map[*WebConn]struct{})
 	}
 	i.byUserId[wc.UserId][wc] = struct{}{}
-	i.byConnection[wc] = connMetadata{
-		channelIDs: channelIDs,
-	}
+	i.byConnection[wc] = channelIDs
 	i.byConnectionId[wc.GetConnectionID()] = wc
 	return nil
 }

 func (i *hubConnectionIndex) Remove(wc *WebConn) {
-	connMeta, ok := i.byConnection[wc]
+	channelIDs, ok := i.byConnection[wc]
 	if !ok {
 		return
 	}
@@ -840,7 +841,7 @@ func (i *hubConnectionIndex) Remove(wc *WebConn) {

 	if i.fastIteration {
 		// Remove from byChannelID for each channel
-		for _, chID := range connMeta.channelIDs {
+		for _, chID := range channelIDs {
 			if channelConns, ok := i.byChannelID[chID]; ok {
 				delete(channelConns, wc)
 			}
@@ -862,10 +863,10 @@ func (i *hubConnectionIndex) InvalidateCMCacheForUser(userID string) error {
 	conns := i.ForUser(userID)

 	// Remove all user connections from existing channels
-	for _, conn := range conns {
-		if meta, ok := i.byConnection[conn]; ok {
+	for conn := range conns {
+		if channelIDs, ok := i.byConnection[conn]; ok {
 			// Remove from old channels
-			for _, chID := range meta.channelIDs {
+			for _, chID := range channelIDs {
 				if channelConns, ok := i.byChannelID[chID]; ok {
 					delete(channelConns, conn)
 				}
@@ -874,7 +875,7 @@ func (i *hubConnectionIndex) InvalidateCMCacheForUser(userID string) error {
 	}

 	// Add connections to new channels
-	for _, conn := range conns {
+	for conn := range conns {
 		newChannelIDs := make([]string, 0, len(cm))
 		for chID := range cm {
 			newChannelIDs = append(newChannelIDs, chID)
@@ -886,9 +887,8 @@ func (i *hubConnectionIndex) InvalidateCMCacheForUser(userID string) error {
 		}

 		// Update connection metadata
-		if meta, ok := i.byConnection[conn]; ok {
-			meta.channelIDs = newChannelIDs
-			i.byConnection[conn] = meta
+		if _, ok := i.byConnection[conn]; ok {
+			i.byConnection[conn] = newChannelIDs
 		}
 	}

@@ -901,39 +901,19 @@ func (i *hubConnectionIndex) Has(wc *WebConn) bool {
 }

 // ForUser returns all connections for a user ID.
-func (i *hubConnectionIndex) ForUser(id string) []*WebConn {
-	userConns, ok := i.byUserId[id]
-	if !ok {
-		return nil
-	}
-
-	// Move to using maps.Keys to use the iterator pattern with 1.23.
-	// This saves the additional slice copy.
-	conns := make([]*WebConn, 0, len(userConns))
-	for conn := range userConns {
-		conns = append(conns, conn)
-	}
-	return conns
+func (i *hubConnectionIndex) ForUser(id string) iter.Seq[*WebConn] {
+	return maps.Keys(i.byUserId[id])
 }

 // ForChannel returns all connections for a channelID.
-func (i *hubConnectionIndex) ForChannel(channelID string) []*WebConn {
-	channelConns, ok := i.byChannelID[channelID]
-	if !ok {
-		return nil
-	}
-
-	conns := make([]*WebConn, 0, len(channelConns))
-	for conn := range channelConns {
-		conns = append(conns, conn)
-	}
-	return conns
+func (i *hubConnectionIndex) ForChannel(channelID string) iter.Seq[*WebConn] {
+	return maps.Keys(i.byChannelID[channelID])
 }

 // 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) {
+	for conn := range i.ForUser(id) {
 		if conn.Active.Load() {
 			cnt++
 		}
@@ -947,7 +927,7 @@ func (i *hubConnectionIndex) ForConnection(id string) *WebConn {
 }

 // All returns the full webConn index.
-func (i *hubConnectionIndex) All() map[*WebConn]connMetadata {
+func (i *hubConnectionIndex) All() map[*WebConn][]string {
 	return i.byConnection
 }

@@ -958,7 +938,7 @@ func (i *hubConnectionIndex) RemoveInactiveByConnectionID(userID, connectionID s
 	if userID == "" {
 		return nil
 	}
-	for _, conn := range i.ForUser(userID) {
+	for conn := range i.ForUser(userID) {
 		if conn.GetConnectionID() == connectionID && !conn.Active.Load() {
 			i.Remove(conn)
 			return conn