From 509b8e9af770f372777a3295cf3fc52f92db99bf Mon Sep 17 00:00:00 2001 From: Agniva De Sarker Date: Sun, 11 May 2025 12:00:12 +0530 Subject: [PATCH] MM-63130: Move to webHub iteration to be alloc-free (#30792) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- server/channels/app/platform/web_hub.go | 108 +++++------ server/channels/app/platform/web_hub_test.go | 178 +++++++++++++++++-- 2 files changed, 204 insertions(+), 82 deletions(-) diff --git a/server/channels/app/platform/web_hub.go b/server/channels/app/platform/web_hub.go index 5eab747a1c..d88fe19501 100644 --- 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 diff --git a/server/channels/app/platform/web_hub_test.go b/server/channels/app/platform/web_hub_test.go index 532ebf3156..dc6ad09264 100644 --- a/server/channels/app/platform/web_hub_test.go +++ b/server/channels/app/platform/web_hub_test.go @@ -7,10 +7,12 @@ import ( "bytes" "encoding/json" "fmt" + "iter" "net" "net/http" "net/http/httptest" "runtime" + "slices" "testing" "time" @@ -254,8 +256,8 @@ func TestHubConnIndex(t *testing.T) { assert.True(t, connIndex.Has(wc1)) assert.True(t, connIndex.Has(wc2)) - assert.ElementsMatch(t, connIndex.ForUser(wc2.UserId), []*WebConn{wc2, wc3, wc4}) - assert.ElementsMatch(t, connIndex.ForUser(wc1.UserId), []*WebConn{wc1}) + assert.ElementsMatch(t, slices.Collect(connIndex.ForUser(wc2.UserId)), []*WebConn{wc2, wc3, wc4}) + assert.ElementsMatch(t, slices.Collect(connIndex.ForUser(wc1.UserId)), []*WebConn{wc1}) assert.True(t, connIndex.Has(wc2)) assert.True(t, connIndex.Has(wc1)) assert.Len(t, connIndex.All(), 4) @@ -264,8 +266,8 @@ func TestHubConnIndex(t *testing.T) { t.Run("RemoveMiddleUser2", func(t *testing.T) { connIndex.Remove(wc3) // Remove from middle from user2 - assert.ElementsMatch(t, connIndex.ForUser(wc2.UserId), []*WebConn{wc2, wc4}) - assert.ElementsMatch(t, connIndex.ForUser(wc1.UserId), []*WebConn{wc1}) + assert.ElementsMatch(t, slices.Collect(connIndex.ForUser(wc2.UserId)), []*WebConn{wc2, wc4}) + assert.ElementsMatch(t, slices.Collect(connIndex.ForUser(wc1.UserId)), []*WebConn{wc1}) assert.True(t, connIndex.Has(wc2)) assert.False(t, connIndex.Has(wc3)) assert.True(t, connIndex.Has(wc4)) @@ -275,9 +277,9 @@ func TestHubConnIndex(t *testing.T) { t.Run("RemoveUser1", func(t *testing.T) { connIndex.Remove(wc1) // Remove sole connection from user1 - 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.ElementsMatch(t, slices.Collect(connIndex.ForUser(wc2.UserId)), []*WebConn{wc2, wc4}) + assert.ElementsMatch(t, slices.Collect(connIndex.ForUser(wc1.UserId)), []*WebConn{}) + assert.Len(t, slices.Collect(connIndex.ForUser(wc1.UserId)), 0) assert.Len(t, connIndex.All(), 2) assert.False(t, connIndex.Has(wc1)) assert.True(t, connIndex.Has(wc2)) @@ -286,8 +288,8 @@ func TestHubConnIndex(t *testing.T) { t.Run("RemoveEndUser2", func(t *testing.T) { connIndex.Remove(wc4) // Remove from end from user2 - assert.ElementsMatch(t, connIndex.ForUser(wc2.UserId), []*WebConn{wc2}) - assert.ElementsMatch(t, connIndex.ForUser(wc1.UserId), []*WebConn{}) + assert.ElementsMatch(t, slices.Collect(connIndex.ForUser(wc2.UserId)), []*WebConn{wc2}) + assert.ElementsMatch(t, slices.Collect(connIndex.ForUser(wc1.UserId)), []*WebConn{}) assert.True(t, connIndex.Has(wc2)) assert.False(t, connIndex.Has(wc3)) assert.False(t, connIndex.Has(wc4)) @@ -400,11 +402,11 @@ func TestHubConnIndex(t *testing.T) { t.Run("ForChannel", func(t *testing.T) { require.Len(t, connIndex.byChannelID, 1) ids := make([]string, 0) - for _, c := range connIndex.ForChannel(th.BasicChannel.Id) { + for c := range connIndex.ForChannel(th.BasicChannel.Id) { ids = append(ids, c.GetConnectionID()) } require.ElementsMatch(t, []string{wc1ID, wc2ID, wc3ID}, ids) - require.Len(t, connIndex.ForChannel("notexist"), 0) + require.Len(t, slices.Collect(connIndex.ForChannel("notexist")), 0) }) ch := th.CreateChannel(th.BasicTeam) @@ -420,14 +422,14 @@ func TestHubConnIndex(t *testing.T) { t.Run("InvalidateCMCacheForUser", func(t *testing.T) { require.NoError(t, connIndex.InvalidateCMCacheForUser(th.BasicUser2.Id)) require.Len(t, connIndex.byChannelID, 2) - require.Len(t, connIndex.ForChannel(th.BasicChannel.Id), 3) - require.Len(t, connIndex.ForChannel(ch.Id), 2) + require.Len(t, slices.Collect(connIndex.ForChannel(th.BasicChannel.Id)), 3) + require.Len(t, slices.Collect(connIndex.ForChannel(ch.Id)), 2) }) t.Run("Remove", func(t *testing.T) { connIndex.Remove(wc3) require.Len(t, connIndex.byChannelID, 2) - require.Len(t, connIndex.ForChannel(th.BasicChannel.Id), 2) + require.Len(t, slices.Collect(connIndex.ForChannel(th.BasicChannel.Id)), 2) }) }) } @@ -470,7 +472,7 @@ func TestHubConnIndexIncorrectRemoval(t *testing.T) { err = connIndex.Add(wc4) require.NoError(t, err) - for _, wc := range connIndex.ForUser(wc2.UserId) { + for wc := range connIndex.ForUser(wc2.UserId) { if !connIndex.Has(wc) { require.Failf(t, "Failed to find connection", "connection: %v", wc) continue @@ -527,21 +529,21 @@ func TestHubConnIndexInactive(t *testing.T) { 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) + assert.Len(t, slices.Collect(connIndex.ForUser(wc2.UserId)), 1) wc3.lastUserActivityAt = model.GetMillis() err = connIndex.Add(wc3) require.NoError(t, err) connIndex.RemoveInactiveConnections() assert.True(t, connIndex.Has(wc3)) - assert.Len(t, connIndex.ForUser(wc2.UserId), 2) + assert.Len(t, slices.Collect(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.Len(t, slices.Collect(connIndex.ForUser(wc2.UserId)), 1) assert.Equal(t, connIndex.ForUserActiveCount(wc2.UserId), 1) assert.Len(t, connIndex.All(), 2) } @@ -647,6 +649,146 @@ func TestHubWebConnCount(t *testing.T) { assert.Equal(t, 0, th.Service.WebConnCountForUser("none")) } +var globalIter iter.Seq[*WebConn] + +func BenchmarkHubConnIndexIteratorForUser(b *testing.B) { + th := Setup(b) + defer th.TearDown() + + connIndex := newHubConnectionIndex(2*time.Second, th.Service.Store, th.Service.logger, false) + + // User1 + wc1 := &WebConn{ + Platform: th.Service, + UserId: model.NewId(), + } + wc1.Active.Store(true) + wc1.SetConnectionID("conn1") + wc1.SetSession(&model.Session{}) + + // User2 + wc2 := &WebConn{ + Platform: th.Service, + UserId: model.NewId(), + } + wc2.Active.Store(true) + wc2.SetConnectionID("conn2") + wc2.SetSession(&model.Session{}) + + wc3 := &WebConn{ + Platform: th.Service, + UserId: wc2.UserId, + } + wc3.Active.Store(false) + wc3.SetConnectionID("conn3") + wc3.SetSession(&model.Session{}) + + require.NoError(b, connIndex.Add(wc1)) + require.NoError(b, connIndex.Add(wc2)) + require.NoError(b, connIndex.Add(wc3)) + + b.ResetTimer() + b.Run("2 users", func(b *testing.B) { + for i := 0; i < b.N; i++ { + globalIter = connIndex.ForUser(wc2.UserId) + } + }) + + wc4 := &WebConn{ + Platform: th.Service, + UserId: wc2.UserId, + } + wc4.Active.Store(false) + wc4.SetConnectionID("conn4") + wc4.SetSession(&model.Session{}) + + require.NoError(b, connIndex.Add(wc4)) + b.ResetTimer() + b.Run("3 users", func(b *testing.B) { + for i := 0; i < b.N; i++ { + globalIter = connIndex.ForUser(wc2.UserId) + } + }) + + wc5 := &WebConn{ + Platform: th.Service, + UserId: wc2.UserId, + } + wc5.Active.Store(false) + wc5.SetConnectionID("conn5") + wc5.SetSession(&model.Session{}) + + require.NoError(b, connIndex.Add(wc5)) + b.ResetTimer() + b.Run("4 users", func(b *testing.B) { + for i := 0; i < b.N; i++ { + globalIter = connIndex.ForUser(wc2.UserId) + } + }) +} + +func BenchmarkHubConnIndexIteratorForChannel(b *testing.B) { + th := Setup(b).InitBasic() + defer th.TearDown() + + _, err := th.Service.Store.Channel().SaveMember(th.Context, &model.ChannelMember{ + ChannelId: th.BasicChannel.Id, + UserId: th.BasicUser.Id, + NotifyProps: model.GetDefaultChannelNotifyProps(), + SchemeGuest: th.BasicUser.IsGuest(), + SchemeUser: !th.BasicUser.IsGuest(), + }) + require.NoError(b, err) + _, err = th.Service.Store.Channel().SaveMember(th.Context, &model.ChannelMember{ + ChannelId: th.BasicChannel.Id, + UserId: th.BasicUser2.Id, + NotifyProps: model.GetDefaultChannelNotifyProps(), + SchemeGuest: th.BasicUser2.IsGuest(), + SchemeUser: !th.BasicUser2.IsGuest(), + }) + require.NoError(b, err) + + connIndex := newHubConnectionIndex(1*time.Second, th.Service.Store, th.Service.logger, true) + + // User1 + wc1ID := model.NewId() + wc1 := &WebConn{ + Platform: th.Service, + Suite: th.Suite, + UserId: th.BasicUser.Id, + } + wc1.SetConnectionID(wc1ID) + wc1.SetSession(&model.Session{}) + + // User2 + wc2ID := model.NewId() + wc2 := &WebConn{ + Platform: th.Service, + Suite: th.Suite, + UserId: th.BasicUser2.Id, + } + wc2.SetConnectionID(wc2ID) + wc2.SetSession(&model.Session{}) + + wc3ID := model.NewId() + wc3 := &WebConn{ + Platform: th.Service, + Suite: th.Suite, + UserId: wc2.UserId, + } + wc3.SetConnectionID(wc3ID) + wc3.SetSession(&model.Session{}) + + require.NoError(b, connIndex.Add(wc1)) + require.NoError(b, connIndex.Add(wc2)) + require.NoError(b, connIndex.Add(wc3)) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + globalIter = connIndex.ForChannel(th.BasicChannel.Id) + } +} + // Always run this with -benchtime=0.1s // See: https://github.com/golang/go/issues/27217. func BenchmarkHubConnIndex(b *testing.B) {