From 45b025898faeec98569b5900c5adc556f9adcbcf Mon Sep 17 00:00:00 2001
From: Ganesh Vernekar <ganeshvern@gmail.com>
Date: Mon, 20 Feb 2023 13:58:06 +0530
Subject: [PATCH] Add BenchmarkHeadChunkQuerier and BenchmarkHeadQuerier

Signed-off-by: Ganesh Vernekar <ganeshvern@gmail.com>
---
 tsdb/querier_test.go | 90 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 90 insertions(+)

diff --git a/tsdb/querier_test.go b/tsdb/querier_test.go
index e6e9f143f..ae9ddad01 100644
--- a/tsdb/querier_test.go
+++ b/tsdb/querier_test.go
@@ -2246,3 +2246,93 @@ func TestBlockBaseSeriesSet(t *testing.T) {
 		require.NoError(t, bcs.Err())
 	}
 }
+
+func BenchmarkHeadChunkQuerier(b *testing.B) {
+	db := openTestDB(b, nil, nil)
+	defer func() {
+		require.NoError(b, db.Close())
+	}()
+
+	// 3h of data.
+	numTimeseries := 100
+	app := db.Appender(context.Background())
+	for i := 0; i < 120*6; i++ {
+		for j := 0; j < numTimeseries; j++ {
+			lbls := labels.FromStrings("foo", fmt.Sprintf("bar%d", j))
+			if i%10 == 0 {
+				require.NoError(b, app.Commit())
+				app = db.Appender(context.Background())
+			}
+			_, err := app.Append(0, lbls, int64(i*15)*time.Second.Milliseconds(), float64(i*100))
+			require.NoError(b, err)
+		}
+	}
+	require.NoError(b, app.Commit())
+
+	querier, err := db.ChunkQuerier(context.Background(), math.MinInt64, math.MaxInt64)
+	require.NoError(b, err)
+	defer func(q storage.ChunkQuerier) {
+		require.NoError(b, q.Close())
+	}(querier)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ss := querier.Select(false, nil, labels.MustNewMatcher(labels.MatchRegexp, "foo", "bar.*"))
+		total := 0
+		for ss.Next() {
+			cs := ss.At()
+			it := cs.Iterator(nil)
+			for it.Next() {
+				m := it.At()
+				total += m.Chunk.NumSamples()
+			}
+		}
+		_ = total
+		require.NoError(b, ss.Err())
+	}
+}
+
+func BenchmarkHeadQuerier(b *testing.B) {
+	db := openTestDB(b, nil, nil)
+	defer func() {
+		require.NoError(b, db.Close())
+	}()
+
+	// 3h of data.
+	numTimeseries := 100
+	app := db.Appender(context.Background())
+	for i := 0; i < 120*6; i++ {
+		for j := 0; j < numTimeseries; j++ {
+			lbls := labels.FromStrings("foo", fmt.Sprintf("bar%d", j))
+			if i%10 == 0 {
+				require.NoError(b, app.Commit())
+				app = db.Appender(context.Background())
+			}
+			_, err := app.Append(0, lbls, int64(i*15)*time.Second.Milliseconds(), float64(i*100))
+			require.NoError(b, err)
+		}
+	}
+	require.NoError(b, app.Commit())
+
+	querier, err := db.Querier(context.Background(), math.MinInt64, math.MaxInt64)
+	require.NoError(b, err)
+	defer func(q storage.Querier) {
+		require.NoError(b, q.Close())
+	}(querier)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ss := querier.Select(false, nil, labels.MustNewMatcher(labels.MatchRegexp, "foo", "bar.*"))
+		total := int64(0)
+		for ss.Next() {
+			cs := ss.At()
+			it := cs.Iterator(nil)
+			for it.Next() != chunkenc.ValNone {
+				ts, _ := it.At()
+				total += ts
+			}
+		}
+		_ = total
+		require.NoError(b, ss.Err())
+	}
+}