From 67f961e4dcc66a9bd84654b701932590cc21eee2 Mon Sep 17 00:00:00 2001
From: Owen Diehl <ow.diehl@gmail.com>
Date: Wed, 9 Mar 2022 09:57:47 -0500
Subject: [PATCH] more tsdb benchmarks + updates script to work with
 incompatible branches (#5590)

* more tsdb benchmarks + updates script to work with incompatible branches

* shellcheck
---
 tools/tsdb/tsdb-map/diff.sh      | 30 +++++++++++-----
 tools/tsdb/tsdb-map/main_test.go | 60 ++++++++++++++++++++++++--------
 2 files changed, 68 insertions(+), 22 deletions(-)

diff --git a/tools/tsdb/tsdb-map/diff.sh b/tools/tsdb/tsdb-map/diff.sh
index f676f33ee1..47c5714e21 100755
--- a/tools/tsdb/tsdb-map/diff.sh
+++ b/tools/tsdb/tsdb-map/diff.sh
@@ -1,18 +1,32 @@
 #!/usr/bin/env bash
 
-old=$1
-new=$2
+# This can be run like:
+# ./tools/tsdb/tsdb-map/diff.sh /tmp/loki-scratch/loki-ops-daily.r main $(git rev-parse --abbrev-ref HEAD)
+
+boltdb_base=$1
+branch_a=$2
+branch_b=$3
+
+echo building from "${branch_a}"
+git checkout "${branch_a}"
+go run tools/tsdb/tsdb-map/main.go  -source "${boltdb_base}" -dest /tmp/loki-tsdb-a
+echo benchmarking "${branch_a}"
+LOKI_TSDB_PATH=/tmp/loki-tsdb-a go test github.com/grafana/loki/tools/tsdb/tsdb-map -bench=BenchmarkQuery -run '^$' -benchmem > /tmp/loki-tsdb-bench-a
+
+echo building from "${branch_b}"
+git checkout "${branch_b}"
+go run tools/tsdb/tsdb-map/main.go  -source "${boltdb_base}" -dest /tmp/loki-tsdb-b
+echo benchmarking "${branch_b}"
+LOKI_TSDB_PATH=/tmp/loki-tsdb-b go test github.com/grafana/loki/tools/tsdb/tsdb-map -bench=BenchmarkQuery -run '^$' -benchmem > /tmp/loki-tsdb-bench-b
+
 
 echo benchmarks:
 echo
-
-benchstat \
-<(LOKI_TSDB_PATH="${old}" go test github.com/grafana/loki/tools/tsdb/tsdb-map -bench=BenchmarkQuery -run '^$' -benchmem) \
-<(LOKI_TSDB_PATH="${new}" go test github.com/grafana/loki/tools/tsdb/tsdb-map -bench=BenchmarkQuery -run '^$' -benchmem)
+benchstat /tmp/loki-tsdb-bench-a /tmp/loki-tsdb-bench-b
 
 echo
 echo sizing:
 echo
 
-ls -lh "${old}"
-ls -lh "${new}"
+ls -lh /tmp/loki-tsdb-a
+ls -lh /tmp/loki-tsdb-b
diff --git a/tools/tsdb/tsdb-map/main_test.go b/tools/tsdb/tsdb-map/main_test.go
index f3c432dd16..e11764df1f 100644
--- a/tools/tsdb/tsdb-map/main_test.go
+++ b/tools/tsdb/tsdb-map/main_test.go
@@ -1,7 +1,9 @@
 package main
 
 import (
+	"context"
 	"fmt"
+	"math"
 	"math/rand"
 	"os"
 	"testing"
@@ -19,21 +21,26 @@ func TestExtractChecksum(t *testing.T) {
 	require.Equal(t, x, extractChecksumFromChunkID([]byte(s)))
 }
 
+type testCase struct {
+	name     string
+	matchers []*labels.Matcher
+}
+
+var cases = []testCase{
+	{
+		name:     "match ns",
+		matchers: []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "namespace", "loki-ops")},
+	},
+	{
+		name:     "match ns regexp",
+		matchers: []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "namespace", "loki-ops")},
+	},
+}
+
+// Only iterates through the low level IndexReader using `PostingsForMatchers`
 // Requires LOKI_TSDB_PATH to be set or this will short-circuit
-func BenchmarkQuery(b *testing.B) {
-	for _, bm := range []struct {
-		name     string
-		matchers []*labels.Matcher
-	}{
-		{
-			name:     "match ns",
-			matchers: []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "namespace", "loki-ops")},
-		},
-		{
-			name:     "match ns regexp",
-			matchers: []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "namespace", "loki-ops")},
-		},
-	} {
+func BenchmarkQuery_PostingsForMatchers(b *testing.B) {
+	for _, bm := range cases {
 		indexPath := os.Getenv("LOKI_TSDB_PATH")
 		if indexPath == "" {
 			return
@@ -53,3 +60,28 @@ func BenchmarkQuery(b *testing.B) {
 		})
 	}
 }
+
+// Uses the higher level loki index interface, resolving chunk refs.
+// Requires LOKI_TSDB_PATH to be set or this will short-circuit
+func BenchmarkQuery_GetChunkRefs(b *testing.B) {
+	for _, bm := range cases {
+		indexPath := os.Getenv("LOKI_TSDB_PATH")
+		if indexPath == "" {
+			return
+		}
+
+		reader, err := index.NewFileReader(indexPath)
+		if err != nil {
+			panic(err)
+		}
+		idx := tsdb.NewTSDBIndex(reader)
+		b.Run(bm.name, func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				_, err := idx.GetChunkRefs(context.Background(), "fake", 0, math.MaxInt64, bm.matchers...)
+				if err != nil {
+					panic(err)
+				}
+			}
+		})
+	}
+}