bucket_bench_test.go

package store

import (
	"context"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/fortytw2/leaktest"
	"github.com/go-kit/kit/log"
	"github.com/improbable-eng/thanos/pkg/block"
	"github.com/improbable-eng/thanos/pkg/objstore"
	"github.com/improbable-eng/thanos/pkg/objstore/inmem"
	"github.com/improbable-eng/thanos/pkg/runutil"
	"github.com/improbable-eng/thanos/pkg/store/storepb"
	"github.com/improbable-eng/thanos/pkg/testutil"
	"github.com/pkg/errors"
	"github.com/prometheus/prometheus/pkg/timestamp"
	"github.com/prometheus/tsdb/labels"
)

func prepareStoreWithTestBlocks2(t testing.TB, ctx context.Context, dir string, bkt objstore.Bucket) (store *BucketStore, minTime, maxTime int64) {
	series := []labels.Labels{
		labels.FromStrings("a", "1", "b", "1"),
		labels.FromStrings("a", "1", "b", "2"),
		labels.FromStrings("a", "2", "b", "1"),
		labels.FromStrings("a", "2", "b", "2"),
		labels.FromStrings("a", "1", "c", "1"),
		labels.FromStrings("a", "1", "c", "2"),
		labels.FromStrings("a", "2", "c", "1"),
		labels.FromStrings("a", "2", "c", "2"),
	}
	extLset := labels.FromStrings("ext1", "value1")

	start := time.Now()
	now := start

	minTime = int64(0)
	maxTime = int64(0)
	for i := 0; i < 3; i++ {
		mint := timestamp.FromTime(now)
		now = now.Add(2 * time.Hour)
		maxt := timestamp.FromTime(now)

		if minTime == 0 {
			minTime = mint
		}
		maxTime = maxt

		// Create two blocks per time slot. Only add 10 samples each so only one chunk
		// gets created each. This way we can easily verify we got 10 chunks per series below.
		id1, err := testutil.CreateBlock(dir, series[:4], 10, mint, maxt, extLset, 0)
		testutil.Ok(t, err)
		id2, err := testutil.CreateBlock(dir, series[4:], 10, mint, maxt, extLset, 0)
		testutil.Ok(t, err)

		dir1, dir2 := filepath.Join(dir, id1.String()), filepath.Join(dir, id2.String())

		// Add labels to the meta of the second block.
		meta, err := block.ReadMetaFile(dir2)
		testutil.Ok(t, err)
		meta.Thanos.Labels = map[string]string{"ext2": "value2"}
		testutil.Ok(t, block.WriteMetaFile(log.NewNopLogger(), dir2, meta))

		testutil.Ok(t, block.Upload(ctx, log.NewNopLogger(), bkt, dir1))
		testutil.Ok(t, block.Upload(ctx, log.NewNopLogger(), bkt, dir2))

		testutil.Ok(t, os.RemoveAll(dir1))
		testutil.Ok(t, os.RemoveAll(dir2))
	}

	store, err := NewBucketStore(nil, nil, bkt, dir, 100, 0, false)
	testutil.Ok(t, err)

	go func() {
		if err := runutil.Repeat(100*time.Millisecond, ctx.Done(), func() error {
			return store.SyncBlocks(ctx)
		}); err != nil && errors.Cause(err) != context.Canceled {
			t.Error(err)
			t.FailNow()
		}
	}()

	rctx, rcancel := context.WithTimeout(ctx, 30*time.Second)
	defer rcancel()
	testutil.Ok(t, runutil.Retry(100*time.Millisecond, rctx.Done(), func() error {
		if store.numBlocks() < 6 {
			return errors.New("not all blocks loaded")
		}
		return nil
	}))

	return store, minTime, maxTime
}

func BenchmarkBucketStore_Series(b *testing.B) {
	defer leaktest.CheckTimeout(b, 10*time.Second)()

	bkt := inmem.NewBucket()

	dir, err := ioutil.TempDir("", "test_bucketstore_bench_e2e")
	testutil.Ok(b, err)
	defer func() { testutil.Ok(b, os.RemoveAll(dir)) }()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	store, minTime, maxTime := prepareStoreWithTestBlocks2(b, ctx, dir, bkt)

	req := &storepb.SeriesRequest{
		Matchers: []storepb.LabelMatcher{
			{Type: storepb.LabelMatcher_RE, Name: "a", Value: "1|2"},
		},
		MinTime: minTime,
		MaxTime: maxTime,
	}

	srv := newStoreSeriesServer(ctx)
	for i := 0; i < b.N; i++ {
		testutil.Ok(b, store.Series(req, srv))
		srv.SeriesSet = []storepb.Series(nil)
	}
}