Optimize PrevEventIDs when getting thousands of backwards extremeties

federationapi/routing/backfill.go

@@ -95,6 +95,12 @@ func Backfill(
 		}
 	}
 
+	// Enforce a limit of 100 events, as not to hit the DB to hard.
+	// Synapse has a hard limit of 100 events as well.
+	if req.Limit > 100 {
+		req.Limit = 100
+	}
+
 	// Query the Roomserver.
 	if err = rsAPI.PerformBackfill(httpReq.Context(), &req, &res); err != nil {
 		util.GetLogger(httpReq.Context()).WithError(err).Error("query.PerformBackfill failed")

roomserver/api/perform.go

@@ -8,7 +8,6 @@ import (
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/gomatrixserverlib/spec"
-	"github.com/matrix-org/util"
 )
 
 type PerformCreateRoomRequest struct {
@@ -91,14 +90,44 @@ type PerformBackfillRequest struct {
 	VirtualHost spec.ServerName `json:"virtual_host"`
 }
 
-// PrevEventIDs returns the prev_event IDs of all backwards extremities, de-duplicated in a lexicographically sorted order.
+// limitPrevEventIDs is the maximum of eventIDs we
+// return when calling PrevEventIDs.
+const limitPrevEventIDs = 100
+
+// PrevEventIDs returns the prev_event IDs of either 100 backwards extremities or
+// len(r.BackwardsExtremities). Limited to 100, due to Synapse/Dendrite stopping after reaching
+// this limit. (which sounds sane)
 func (r *PerformBackfillRequest) PrevEventIDs() []string {
-	var prevEventIDs []string
+	var uniqueIDs map[string]struct{}
+
+	// Create a unique eventID map of either 100 or len(r.BackwardsExtremities).
+	// 100 since Synapse/Dendrite stops after reaching 100 events.
+	if len(r.BackwardsExtremities) > limitPrevEventIDs {
+		uniqueIDs = make(map[string]struct{}, limitPrevEventIDs)
+	} else {
+		uniqueIDs = make(map[string]struct{}, len(r.BackwardsExtremities))
+	}
+
+outerLoop:
 	for _, pes := range r.BackwardsExtremities {
-		prevEventIDs = append(prevEventIDs, pes...)
+		for _, evID := range pes {
+			uniqueIDs[evID] = struct{}{}
+			// We found enough unique eventIDs.
+			if len(uniqueIDs) >= limitPrevEventIDs {
+				break outerLoop
+			}
+		}
 	}
-	prevEventIDs = util.UniqueStrings(prevEventIDs)
-	return prevEventIDs
+
+	// map -> []string
+	result := make([]string, len(uniqueIDs))
+	i := 0
+	for evID := range uniqueIDs {
+		result[i] = evID
+		i++
+	}
+
+	return result
 }
 
 // PerformBackfillResponse is a response to PerformBackfill.

roomserver/api/perform_test.go

@@ -0,0 +1,81 @@
+package api
+
+import (
+	"fmt"
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func BenchmarkPrevEventIDs(b *testing.B) {
+	for _, x := range []int64{1, 10, 100, 500, 1000, 2000} {
+		benchPrevEventIDs(b, int(x))
+	}
+}
+
+func benchPrevEventIDs(b *testing.B, count int) {
+	bwExtrems := generateBackwardsExtremities(b, count)
+	backfiller := PerformBackfillRequest{
+		BackwardsExtremities: bwExtrems,
+	}
+
+	b.Run(fmt.Sprintf("Original%d", count), func(b *testing.B) {
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			prevIDs := backfiller.PrevEventIDs()
+			_ = prevIDs
+		}
+	})
+}
+
+type testLike interface {
+	Helper()
+}
+
+const randomIDCharsCount = 10
+
+func generateBackwardsExtremities(t testLike, count int) map[string][]string {
+	t.Helper()
+	result := make(map[string][]string, count)
+	for i := 0; i < count; i++ {
+		eventID := randomEventId(int64(i))
+		result[eventID] = []string{
+			randomEventId(int64(i + 1)),
+			randomEventId(int64(i + 2)),
+			randomEventId(int64(i + 3)),
+		}
+	}
+	return result
+}
+
+const alphanumerics = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+
+// randomEventId generates a pseudo-random string of length n.
+func randomEventId(src int64) string {
+	randSrc := rand.NewSource(src)
+	b := make([]byte, randomIDCharsCount)
+	for i := range b {
+		b[i] = alphanumerics[randSrc.Int63()%int64(len(alphanumerics))]
+	}
+	return string(b)
+}
+
+func TestPrevEventIDs(t *testing.T) {
+	// generate 10 backwards extremities
+	bwExtrems := generateBackwardsExtremities(t, 10)
+	backfiller := PerformBackfillRequest{
+		BackwardsExtremities: bwExtrems,
+	}
+
+	prevIDs := backfiller.PrevEventIDs()
+	// Given how "generateBackwardsExtremities" works, this
+	// generates 12 unique event IDs
+	assert.Equal(t, 12, len(prevIDs))
+
+	// generate 200 backwards extremities
+	backfiller.BackwardsExtremities = generateBackwardsExtremities(t, 200)
+	prevIDs = backfiller.PrevEventIDs()
+	// PrevEventIDs returns at max 100 event IDs
+	assert.Equal(t, 100, len(prevIDs))
+}