https://github.com/thanos-io/thanos/pull/5588

bucket_bench_test.go

package store

import (
	"context"
	"math/rand"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/go-kit/log"
	"github.com/oklog/ulid"
	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/tsdb"
	"github.com/thanos-io/objstore"
	"github.com/thanos-io/objstore/providers/filesystem"

	"github.com/thanos-io/thanos/pkg/block"
	"github.com/thanos-io/thanos/pkg/block/indexheader"
	"github.com/thanos-io/thanos/pkg/block/metadata"
	storetestutil "github.com/thanos-io/thanos/pkg/store/storepb/testutil"
	"github.com/thanos-io/thanos/pkg/testutil"
)

func BenchmarkBucketIndexReader_ExpandedPostings2(b *testing.B) {
	tb := testutil.NewTB(b)

	tmpDir := b.TempDir()

	bkt, err := filesystem.NewBucket(filepath.Join(tmpDir, "bkt"))
	testutil.Ok(tb, err)
	defer func() { testutil.Ok(tb, bkt.Close()) }()

	id := uploadTestBlock2(tb, tmpDir, bkt, 50000)
	r, err := indexheader.NewBinaryReader(context.Background(), log.NewNopLogger(), bkt, tmpDir, id, DefaultPostingOffsetInMemorySampling)
	testutil.Ok(tb, err)

	benchmarkExpandedPostings2(tb, bkt, id, r, 50000)
}

func uploadTestBlock2(t testing.TB, tmpDir string, bkt objstore.Bucket, series int) ulid.ULID {
	headOpts := tsdb.DefaultHeadOptions()
	headOpts.ChunkDirRoot = tmpDir
	headOpts.ChunkRange = 1000
	h, err := tsdb.NewHead(nil, nil, nil, headOpts, nil)
	testutil.Ok(t, err)
	defer func() {
		testutil.Ok(t, h.Close())
	}()

	logger := log.NewNopLogger()

	appendTestData2(t, h.Appender(context.Background()), series)

	testutil.Ok(t, os.MkdirAll(filepath.Join(tmpDir, "tmp"), os.ModePerm))
	id := createBlockFromHead(t, filepath.Join(tmpDir, "tmp"), h)

	_, err = metadata.InjectThanos(log.NewNopLogger(), filepath.Join(tmpDir, "tmp", id.String()), metadata.Thanos{
		Labels:     labels.Labels{{Name: "ext1", Value: "1"}}.Map(),
		Downsample: metadata.ThanosDownsample{Resolution: 0},
		Source:     metadata.TestSource,
	}, nil)
	testutil.Ok(t, err)
	testutil.Ok(t, block.Upload(context.Background(), logger, bkt, filepath.Join(tmpDir, "tmp", id.String()), metadata.NoneFunc))
	testutil.Ok(t, block.Upload(context.Background(), logger, bkt, filepath.Join(tmpDir, "tmp", id.String()), metadata.NoneFunc))

	return id
}

func appendTestData2(t testing.TB, app storage.Appender, series int) {
	addSeries := func(l labels.Labels) {
		_, err := app.Append(0, l, 0, 0)
		testutil.Ok(t, err)
	}

	series = series / 5
	for n := 0; n < 10; n++ {
		for i := 0; i < series/10; i++ {
			j := series/10*n + i
			addSeries(labels.FromStrings("i", strconv.Itoa(i)+storetestutil.LabelLongSuffix, "n", strconv.Itoa(j)+storetestutil.LabelLongSuffix, "j", "foo"))
			// Have some series that won't be matched, to properly test inverted matches.
			addSeries(labels.FromStrings("i", strconv.Itoa(i)+storetestutil.LabelLongSuffix, "n", strconv.Itoa(j)+storetestutil.LabelLongSuffix, "j", "bar"))
			addSeries(labels.FromStrings("i", strconv.Itoa(i)+storetestutil.LabelLongSuffix, "n", "0_"+strconv.Itoa(j)+storetestutil.LabelLongSuffix, "j", "bar"))
			addSeries(labels.FromStrings("i", strconv.Itoa(i)+storetestutil.LabelLongSuffix, "n", "1_"+strconv.Itoa(j)+storetestutil.LabelLongSuffix, "j", "bar"))
			addSeries(labels.FromStrings("i", strconv.Itoa(i)+storetestutil.LabelLongSuffix, "n", "2_"+strconv.Itoa(j)+storetestutil.LabelLongSuffix, "j", "foo"))
		}
	}
	testutil.Ok(t, app.Commit())
}

func benchmarkExpandedPostings2(
	t testutil.TB,
	bkt objstore.BucketReader,
	id ulid.ULID,
	r indexheader.Reader,
	series int,
) {
	series = series / 5
	valuesSet := make([]string, 0, series)
	for n := 0; n < 10; n++ {
		for i := 0; i < series/10; i++ {
			j := series/10*n + i
			valuesSet = append(valuesSet, strconv.Itoa(j)+storetestutil.LabelLongSuffix)
		}
	}
	rand.New(rand.NewSource(time.Now().UnixNano())).Shuffle(len(valuesSet), func(i, j int) {
		valuesSet[i], valuesSet[j] = valuesSet[j], valuesSet[i]
	})
	iRegexSet := labels.MustNewMatcher(labels.MatchRegexp, "n", strings.Join(valuesSet, "|"))

	cases := []struct {
		name     string
		matchers []*labels.Matcher

		expectedLen int
	}{
		{`n=~"101|random-shuffled-values|1"`, []*labels.Matcher{iRegexSet}, series * 2},
	}

	for _, c := range cases {
		t.Run(c.name, func(t testutil.TB) {
			b := &bucketBlock{
				logger:            log.NewNopLogger(),
				metrics:           newBucketStoreMetrics(nil),
				indexHeaderReader: r,
				indexCache:        noopCache{},
				bkt:               bkt,
				meta:              &metadata.Meta{BlockMeta: tsdb.BlockMeta{ULID: id}},
				partitioner:       NewGapBasedPartitioner(PartitionerMaxGapSize),
			}

			indexr := newBucketIndexReader(b)

			t.ResetTimer()
			for i := 0; i < t.N(); i++ {
				p, err := indexr.ExpandedPostings(context.Background(), c.matchers)
				testutil.Ok(t, err)
				testutil.Equals(t, c.expectedLen, len(p))
			}
		})
	}
}