MM-24283: Refactor hub connection index (#14461)

Automatic Merge
2020-05-22 10:39:21 +05:30
--- 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
 }
--- 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()