diff --git a/app/web_hub.go b/app/web_hub.go index d19813c45a..d5deaadb16 100644 --- a/app/web_hub.go +++ b/app/web_hub.go @@ -474,14 +474,10 @@ func (h *Hub) Start() { if metrics := h.app.Metrics(); metrics != nil { metrics.DecrementWebSocketBroadcastBufferSize(strconv.Itoa(h.connectionIndex), 1) } - candidates := connIndex.All() - if msg.GetBroadcast().UserId != "" { - candidates = connIndex.ForUser(msg.GetBroadcast().UserId) - } msg = msg.PrecomputeJSON() - for _, webConn := range candidates { + broadcast := func(webConn *WebConn) { if !connIndex.Has(webConn) { - continue + return } if webConn.shouldSendEvent(msg) { select { @@ -493,8 +489,19 @@ func (h *Hub) Start() { } } } + if msg.GetBroadcast().UserId != "" { + candidates := connIndex.ForUser(msg.GetBroadcast().UserId) + for _, webConn := range candidates { + broadcast(webConn) + } + continue + } + candidates := connIndex.All() + for webConn := range candidates { + broadcast(webConn) + } case <-h.stop: - for _, webConn := range connIndex.All() { + for webConn := range connIndex.All() { webConn.Close() h.app.SetStatusOffline(webConn.UserId, false) } @@ -529,64 +536,60 @@ func (h *Hub) Start() { go doRecoverableStart() } -type hubConnectionIndexIndexes struct { - connections int - connectionsByUserId int -} - // 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. +// - get all connections for a given userID. +// - get all connections. type hubConnectionIndex struct { - connections []*WebConn - connectionsByUserId map[string][]*WebConn - connectionIndexes map[*WebConn]*hubConnectionIndexIndexes + // byUserId stores the list of connections for a given userID + byUserId map[string][]*WebConn + // 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 } func newHubConnectionIndex() *hubConnectionIndex { return &hubConnectionIndex{ - connections: make([]*WebConn, 0, model.SESSION_CACHE_SIZE), - connectionsByUserId: make(map[string][]*WebConn), - connectionIndexes: make(map[*WebConn]*hubConnectionIndexIndexes), + byUserId: make(map[string][]*WebConn), + byConnection: make(map[*WebConn]int), } } func (i *hubConnectionIndex) Add(wc *WebConn) { - i.connections = append(i.connections, wc) - i.connectionsByUserId[wc.UserId] = append(i.connectionsByUserId[wc.UserId], wc) - i.connectionIndexes[wc] = &hubConnectionIndexIndexes{ - connections: len(i.connections) - 1, - connectionsByUserId: len(i.connectionsByUserId[wc.UserId]) - 1, - } + i.byUserId[wc.UserId] = append(i.byUserId[wc.UserId], wc) + i.byConnection[wc] = len(i.byUserId[wc.UserId]) - 1 } func (i *hubConnectionIndex) Remove(wc *WebConn) { - indexes, ok := i.connectionIndexes[wc] + userConnIndex, ok := i.byConnection[wc] if !ok { return } - last := i.connections[len(i.connections)-1] - i.connections[indexes.connections] = last - i.connections = i.connections[:len(i.connections)-1] - i.connectionIndexes[last].connections = indexes.connections + // get the conn slice. + userConnections := i.byUserId[wc.UserId] + // get the last connection. + last := userConnections[len(userConnections)-1] + // set the slot that we are trying to remove to be the last connection. + userConnections[userConnIndex] = last + // remove the last connection from the slice. + i.byUserId[wc.UserId] = userConnections[:len(userConnections)-1] + // set the index of the connection that was moved to the new index. + i.byConnection[last] = userConnIndex - userConnections := i.connectionsByUserId[wc.UserId] - last = userConnections[len(userConnections)-1] - userConnections[indexes.connectionsByUserId] = last - i.connectionsByUserId[wc.UserId] = userConnections[:len(userConnections)-1] - i.connectionIndexes[last].connectionsByUserId = indexes.connectionsByUserId - - delete(i.connectionIndexes, wc) + delete(i.byConnection, wc) } func (i *hubConnectionIndex) Has(wc *WebConn) bool { - _, ok := i.connectionIndexes[wc] + _, ok := i.byConnection[wc] return ok } func (i *hubConnectionIndex) ForUser(id string) []*WebConn { - return i.connectionsByUserId[id] + return i.byUserId[id] } -func (i *hubConnectionIndex) All() []*WebConn { - return i.connections +func (i *hubConnectionIndex) All() map[*WebConn]int { + return i.byConnection } diff --git a/app/web_hub_test.go b/app/web_hub_test.go index a1f586f048..60dbfea599 100644 --- a/app/web_hub_test.go +++ b/app/web_hub_test.go @@ -110,6 +110,125 @@ func TestHubStopRaceCondition(t *testing.T) { } } +func TestHubConnIndex(t *testing.T) { + th := Setup(t).InitBasic() + defer th.TearDown() + + connIndex := newHubConnectionIndex() + + // User1 + wc1 := &WebConn{ + App: th.App, + UserId: model.NewId(), + } + + // User2 + wc2 := &WebConn{ + App: th.App, + UserId: model.NewId(), + } + wc3 := &WebConn{ + App: th.App, + UserId: wc2.UserId, + } + wc4 := &WebConn{ + App: th.App, + UserId: wc2.UserId, + } + + connIndex.Add(wc1) + connIndex.Add(wc2) + connIndex.Add(wc3) + connIndex.Add(wc4) + + t.Run("Basic", func(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.True(t, connIndex.Has(wc2)) + assert.True(t, connIndex.Has(wc1)) + assert.Len(t, connIndex.All(), 4) + }) + + 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.True(t, connIndex.Has(wc2)) + assert.False(t, connIndex.Has(wc3)) + assert.True(t, connIndex.Has(wc4)) + assert.Len(t, connIndex.All(), 3) + }) + + 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.All(), 2) + assert.False(t, connIndex.Has(wc1)) + assert.True(t, connIndex.Has(wc2)) + }) + + t.Run("RemoveEndUser2", func(t *testing.T) { + connIndex.Remove(wc4) // Remove from end from user2 + + assert.ElementsMatch(t, connIndex.ForUser(wc2.UserId), []*WebConn{wc4}) + assert.ElementsMatch(t, connIndex.ForUser(wc1.UserId), []*WebConn{}) + assert.True(t, connIndex.Has(wc2)) + assert.False(t, connIndex.Has(wc3)) + assert.False(t, connIndex.Has(wc4)) + assert.Len(t, connIndex.All(), 1) + }) +} + +// Always run this with -benchtime=0.1s +// See: https://github.com/golang/go/issues/27217. +func BenchmarkHubConnIndex(b *testing.B) { + th := Setup(b).InitBasic() + defer th.TearDown() + connIndex := newHubConnectionIndex() + + // User1 + wc1 := &WebConn{ + App: th.App, + UserId: model.NewId(), + } + + // User2 + wc2 := &WebConn{ + App: th.App, + UserId: model.NewId(), + } + b.ResetTimer() + b.Run("Add", func(b *testing.B) { + for i := 0; i < b.N; i++ { + connIndex.Add(wc1) + connIndex.Add(wc2) + + b.StopTimer() + connIndex.Remove(wc1) + connIndex.Remove(wc2) + b.StartTimer() + } + }) + + b.Run("Remove", func(b *testing.B) { + for i := 0; i < b.N; i++ { + b.StopTimer() + connIndex.Add(wc1) + connIndex.Add(wc2) + b.StartTimer() + + connIndex.Remove(wc1) + connIndex.Remove(wc2) + } + }) +} + func TestHubIsRegistered(t *testing.T) { th := Setup(t).InitBasic() defer th.TearDown()