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

@@ -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

tools/tsdb/tsdb-map/main_test.go

@@ -1,7 +1,9 @@
 package main
 
 import (
+	"context"
 	"fmt"
+	"math"
 	"math/rand"
 	"os"
 	"testing"
@@ -19,12 +21,12 @@ func TestExtractChecksum(t *testing.T) {
 	require.Equal(t, x, extractChecksumFromChunkID([]byte(s)))
 }
 
-// Requires LOKI_TSDB_PATH to be set or this will short-circuit
-func BenchmarkQuery(b *testing.B) {
-	for _, bm := range []struct {
+type testCase struct {
 	name     string
 	matchers []*labels.Matcher
-	}{
+}
+
+var cases = []testCase{
 	{
 		name:     "match ns",
 		matchers: []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "namespace", "loki-ops")},
@@ -33,7 +35,12 @@ func BenchmarkQuery(b *testing.B) {
 		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_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)
+				}
+			}
+		})
+	}
+}