From 5b9b1d31a8833dbe10923670c0b0354dcf797470 Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Sat, 20 Jan 2024 19:55:34 +0100
Subject: [PATCH] Add more benchmarks

---
 roomserver/api/perform_test.go | 96 +++++++++++++++++++++++++++++++---
 1 file changed, 88 insertions(+), 8 deletions(-)

diff --git a/roomserver/api/perform_test.go b/roomserver/api/perform_test.go
index 448dc15ad..46cbd76f5 100644
--- a/roomserver/api/perform_test.go
+++ b/roomserver/api/perform_test.go
@@ -5,24 +5,42 @@ import (
 	"math/rand"
 	"testing"
 
+	"github.com/matrix-org/util"
 	"github.com/stretchr/testify/assert"
 )
 
 func BenchmarkPrevEventIDs(b *testing.B) {
-	for _, x := range []int64{1, 2, 4, 10, 50, 100, 300, 500, 800, 1000, 2000, 3000, 5000, 10000} {
+	for _, x := range []int64{1, 10, 100, 500, 1000, 2000} {
 		benchPrevEventIDs(b, int(x))
 	}
 }
 
 func benchPrevEventIDs(b *testing.B, count int) {
-	b.Run(fmt.Sprintf("Benchmark%d", count), func(b *testing.B) {
-		bwExtrems := generateBackwardsExtremities(b, count)
-		backfiller := PerformBackfillRequest{
-			BackwardsExtremities: bwExtrems,
-		}
+	bwExtrems := generateBackwardsExtremities(b, count)
+	backfiller := testPerformBackfillRequest{
+		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 := backfiller.originalPrevEventIDs()
+			_ = prevIDs
+		}
+	})
+
+	b.Run(fmt.Sprintf("FirstAttempt%d", count), func(b *testing.B) {
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			prevIDs := backfiller.firstAttemptPrevEventIDs()
+			_ = prevIDs
+		}
+	})
+
+	b.Run(fmt.Sprintf("UsingMap%d", count), func(b *testing.B) {
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			prevIDs := backfiller.prevEventIDsUsingMap()
 			_ = prevIDs
 		}
 	})
@@ -48,9 +66,71 @@ func generateBackwardsExtremities(t testLike, count int) map[string][]string {
 	return result
 }
 
+// PerformBackfillRequest is a request to PerformBackfill.
+type testPerformBackfillRequest struct {
+	BackwardsExtremities map[string][]string `json:"backwards_extremities"`
+}
+
+func (r *testPerformBackfillRequest) originalPrevEventIDs() []string {
+	var prevEventIDs []string
+	for _, pes := range r.BackwardsExtremities {
+		prevEventIDs = append(prevEventIDs, pes...)
+	}
+	prevEventIDs = util.UniqueStrings(prevEventIDs)
+	return prevEventIDs
+}
+
+func (r *testPerformBackfillRequest) firstAttemptPrevEventIDs() []string {
+	// Collect 1k eventIDs, if possible, they may be cleared out below
+	maxPrevEventIDs := len(r.BackwardsExtremities) * 3
+	if maxPrevEventIDs > 2000 {
+		maxPrevEventIDs = 2000
+	}
+	prevEventIDs := make([]string, 0, maxPrevEventIDs)
+	for _, pes := range r.BackwardsExtremities {
+		prevEventIDs = append(prevEventIDs, pes...)
+		if len(prevEventIDs) > 1000 {
+			break
+		}
+	}
+	prevEventIDs = util.UniqueStrings(prevEventIDs)
+	// If we still have > 100 eventIDs, only return the first 100
+	if len(prevEventIDs) > 100 {
+		return prevEventIDs[:100]
+	}
+	return prevEventIDs
+}
+
+func (r *testPerformBackfillRequest) prevEventIDsUsingMap() []string {
+	var uniqueIDs map[string]struct{}
+	if len(r.BackwardsExtremities) > 100 {
+		uniqueIDs = make(map[string]struct{}, 100)
+	} else {
+		uniqueIDs = make(map[string]struct{}, len(r.BackwardsExtremities))
+	}
+
+outerLoop:
+	for _, pes := range r.BackwardsExtremities {
+		for _, evID := range pes {
+			uniqueIDs[evID] = struct{}{}
+			if len(uniqueIDs) >= 100 {
+				break outerLoop
+			}
+		}
+	}
+
+	result := make([]string, len(uniqueIDs))
+	i := 0
+	for evID := range uniqueIDs {
+		result[i] = evID
+	}
+
+	return result
+}
+
 const alphanumerics = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
 
-// RandomString generates a pseudo-random string of length n.
+// randomEventId generates a pseudo-random string of length n.
 func randomEventId(src int64) string {
 	randSrc := rand.NewSource(src)
 	b := make([]byte, randomIDCharsCount)