From 161b79c75ea20c3e5e0d15e5384259497099599f Mon Sep 17 00:00:00 2001
From: Callum Styan <callumstyan@gmail.com>
Date: Wed, 18 Sep 2024 10:59:21 -0700
Subject: [PATCH] support v2 proto for BenchmarkSampleSend

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

diff --git a/storage/remote/queue_manager_test.go b/storage/remote/queue_manager_test.go
index 6c61a477f..99fd02306 100644
--- a/storage/remote/queue_manager_test.go
+++ b/storage/remote/queue_manager_test.go
@@ -1326,21 +1326,25 @@ func BenchmarkSampleSend(b *testing.B) {
 	cfg.MaxShards = 20
 
 	// todo: test with new proto type(s)
-	m := newTestQueueManager(b, cfg, mcfg, defaultFlushDeadline, c, config.RemoteWriteProtoMsgV1)
-	m.StoreSeries(series, 0)
+	for _, format := range []config.RemoteWriteProtoMsg{config.RemoteWriteProtoMsgV1, config.RemoteWriteProtoMsgV2} {
+		b.Run(string(format), func(b *testing.B) {
+			m := newTestQueueManager(b, cfg, mcfg, defaultFlushDeadline, c, format)
+			m.StoreSeries(series, 0)
 
-	// These should be received by the client.
-	m.Start()
-	defer m.Stop()
+			// These should be received by the client.
+			m.Start()
+			defer m.Stop()
 
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		m.Append(samples)
-		m.UpdateSeriesSegment(series, i+1) // simulate what wlog.Watcher.garbageCollectSeries does
-		m.SeriesReset(i + 1)
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				m.Append(samples)
+				m.UpdateSeriesSegment(series, i+1) // simulate what wlog.Watcher.garbageCollectSeries does
+				m.SeriesReset(i + 1)
+			}
+			// Do not include shutdown
+			b.StopTimer()
+		})
 	}
-	// Do not include shutdown
-	b.StopTimer()
 }
 
 // Check how long it takes to add N series, including external labels processing.