From c328d5fc8820ec7fe45296cfd4de8036b58f4a3a Mon Sep 17 00:00:00 2001
From: Callum Styan <callumstyan@gmail.com>
Date: Tue, 17 Sep 2024 23:04:10 -0700
Subject: [PATCH] fix rwv2 build write request benchmark, also change how the
 memory usage (#14925)

is reported for these benchmarks to more accurately reflect what's
actually allocated

Signed-off-by: Callum Styan <callumstyan@gmail.com>
---
 storage/remote/queue_manager_test.go | 63 ++++++++++++----------------
 1 file changed, 27 insertions(+), 36 deletions(-)

diff --git a/storage/remote/queue_manager_test.go b/storage/remote/queue_manager_test.go
index 38f3fbea2..6c61a477f 100644
--- a/storage/remote/queue_manager_test.go
+++ b/storage/remote/queue_manager_test.go
@@ -1859,13 +1859,6 @@ func BenchmarkBuildWriteRequest(b *testing.B) {
 		}
 		pBuf := proto.NewBuffer(nil)
 
-		// Warmup buffers
-		for i := 0; i < 10; i++ {
-			populateTimeSeries(batch, seriesBuff, true, true)
-			buildWriteRequest(noopLogger, seriesBuff, nil, pBuf, &buff, nil, "snappy")
-		}
-
-		b.ResetTimer()
 		totalSize := 0
 		for i := 0; i < b.N; i++ {
 			populateTimeSeries(batch, seriesBuff, true, true)
@@ -1897,45 +1890,43 @@ func BenchmarkBuildWriteRequest(b *testing.B) {
 
 func BenchmarkBuildV2WriteRequest(b *testing.B) {
 	noopLogger := log.NewNopLogger()
-	type testcase struct {
-		batch []timeSeries
-	}
-	testCases := []testcase{
-		{createDummyTimeSeries(2)},
-		{createDummyTimeSeries(10)},
-		{createDummyTimeSeries(100)},
-	}
-	for _, tc := range testCases {
+	bench := func(b *testing.B, batch []timeSeries) {
 		symbolTable := writev2.NewSymbolTable()
 		buff := make([]byte, 0)
-		seriesBuff := make([]writev2.TimeSeries, len(tc.batch))
+		seriesBuff := make([]writev2.TimeSeries, len(batch))
 		for i := range seriesBuff {
 			seriesBuff[i].Samples = []writev2.Sample{{}}
 			seriesBuff[i].Exemplars = []writev2.Exemplar{{}}
 		}
 		pBuf := []byte{}
 
-		// Warmup buffers
-		for i := 0; i < 10; i++ {
-			populateV2TimeSeries(&symbolTable, tc.batch, seriesBuff, true, true)
-			buildV2WriteRequest(noopLogger, seriesBuff, symbolTable.Symbols(), &pBuf, &buff, nil, "snappy")
-		}
-
-		b.Run(fmt.Sprintf("%d-instances", len(tc.batch)), func(b *testing.B) {
-			totalSize := 0
-			for j := 0; j < b.N; j++ {
-				populateV2TimeSeries(&symbolTable, tc.batch, seriesBuff, true, true)
-				b.ResetTimer()
-				req, _, _, err := buildV2WriteRequest(noopLogger, seriesBuff, symbolTable.Symbols(), &pBuf, &buff, nil, "snappy")
-				if err != nil {
-					b.Fatal(err)
-				}
-				symbolTable.Reset()
-				totalSize += len(req)
-				b.ReportMetric(float64(totalSize)/float64(b.N), "compressedSize/op")
+		totalSize := 0
+		for i := 0; i < b.N; i++ {
+			populateV2TimeSeries(&symbolTable, batch, seriesBuff, true, true)
+			req, _, _, err := buildV2WriteRequest(noopLogger, seriesBuff, symbolTable.Symbols(), &pBuf, &buff, nil, "snappy")
+			if err != nil {
+				b.Fatal(err)
 			}
-		})
+			totalSize += len(req)
+			b.ReportMetric(float64(totalSize)/float64(b.N), "compressedSize/op")
+		}
 	}
+
+	twoBatch := createDummyTimeSeries(2)
+	tenBatch := createDummyTimeSeries(10)
+	hundredBatch := createDummyTimeSeries(100)
+
+	b.Run("2 instances", func(b *testing.B) {
+		bench(b, twoBatch)
+	})
+
+	b.Run("10 instances", func(b *testing.B) {
+		bench(b, tenBatch)
+	})
+
+	b.Run("1k instances", func(b *testing.B) {
+		bench(b, hundredBatch)
+	})
 }
 
 func TestDropOldTimeSeries(t *testing.T) {