github.com/thanos-io/thanos@v0.32.5/pkg/query/query_bench_test.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package query
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"math"
    10  	"math/rand"
    11  	"path/filepath"
    12  	"testing"
    13  	"time"
    14  
    15  	"github.com/go-kit/log"
    16  	"github.com/prometheus/prometheus/model/labels"
    17  	"github.com/prometheus/prometheus/tsdb/chunkenc"
    18  
    19  	"github.com/efficientgo/core/testutil"
    20  
    21  	"github.com/thanos-io/thanos/pkg/gate"
    22  	"github.com/thanos-io/thanos/pkg/store/labelpb"
    23  	"github.com/thanos-io/thanos/pkg/store/storepb"
    24  	storetestutil "github.com/thanos-io/thanos/pkg/store/storepb/testutil"
    25  )
    26  
    27  // TestQuerySelect benchmarks querier Select method. Note that this is what PromQL is using, but PromQL might invoke
    28  // this many times and within different interval e.g
    29  // TODO(bwplotka): Add benchmarks with PromQL involvement.
    30  func TestQuerySelect(t *testing.T) {
    31  	tb := testutil.NewTB(t)
    32  	storetestutil.RunSeriesInterestingCases(tb, 200e3, 200e3, func(t testutil.TB, samplesPerSeries, series int) {
    33  		benchQuerySelect(t, samplesPerSeries, series, true)
    34  	})
    35  }
    36  
    37  func BenchmarkQuerySelect(b *testing.B) {
    38  	tb := testutil.NewTB(b)
    39  	storetestutil.RunSeriesInterestingCases(tb, 10e6, 10e5, func(t testutil.TB, samplesPerSeries, series int) {
    40  		benchQuerySelect(t, samplesPerSeries, series, true)
    41  	})
    42  }
    43  
    44  func benchQuerySelect(t testutil.TB, totalSamples, totalSeries int, dedup bool) {
    45  	tmpDir := t.TempDir()
    46  
    47  	const numOfReplicas = 2
    48  
    49  	samplesPerSeriesPerReplica := totalSamples / numOfReplicas
    50  	if samplesPerSeriesPerReplica == 0 {
    51  		samplesPerSeriesPerReplica = 1
    52  	}
    53  	seriesPerReplica := totalSeries / numOfReplicas
    54  	if seriesPerReplica == 0 {
    55  		seriesPerReplica = 1
    56  	}
    57  
    58  	random := rand.New(rand.NewSource(120))
    59  	var resps []*storepb.SeriesResponse
    60  	var expectedSeries []labels.Labels
    61  	for j := 0; j < numOfReplicas; j++ {
    62  		// Note 0 argument - this is because we want to have two replicas for the same time duration.
    63  		head, created := storetestutil.CreateHeadWithSeries(t, 0, storetestutil.HeadGenOptions{
    64  			TSDBDir:          filepath.Join(tmpDir, fmt.Sprintf("%d", j)),
    65  			SamplesPerSeries: samplesPerSeriesPerReplica,
    66  			Series:           seriesPerReplica,
    67  			Random:           random,
    68  			PrependLabels:    labels.FromStrings("a_replica", fmt.Sprintf("%d", j)), // a_ prefix so we keep sorted order.
    69  		})
    70  		testutil.Ok(t, head.Close())
    71  		for i := 0; i < len(created); i++ {
    72  			if !dedup || j == 0 {
    73  				lset := labelpb.ZLabelsToPromLabels(created[i].Labels).Copy()
    74  				if dedup {
    75  					lset = lset[1:]
    76  				}
    77  				expectedSeries = append(expectedSeries, lset)
    78  			}
    79  
    80  			resps = append(resps, storepb.NewSeriesResponse(created[i]))
    81  		}
    82  
    83  	}
    84  
    85  	logger := log.NewNopLogger()
    86  	q := newQuerier(
    87  		context.Background(),
    88  		logger,
    89  		math.MinInt64,
    90  		math.MaxInt64,
    91  		[]string{"a_replica"},
    92  		nil,
    93  		newProxyStore(&mockedStoreServer{responses: resps}),
    94  		dedup,
    95  		0,
    96  		false,
    97  		false,
    98  		false,
    99  		gate.NewNoop(),
   100  		10*time.Second,
   101  		nil,
   102  		NoopSeriesStatsReporter,
   103  	)
   104  	testSelect(t, q, expectedSeries)
   105  }
   106  
   107  type mockedStoreServer struct {
   108  	storepb.StoreServer
   109  
   110  	responses []*storepb.SeriesResponse
   111  }
   112  
   113  func (m *mockedStoreServer) Series(_ *storepb.SeriesRequest, server storepb.Store_SeriesServer) error {
   114  	for _, r := range m.responses {
   115  		if err := server.Send(r); err != nil {
   116  			return err
   117  		}
   118  	}
   119  	return nil
   120  }
   121  
   122  var (
   123  	testT    int64
   124  	testV    float64
   125  	testLset labels.Labels
   126  )
   127  
   128  func testSelect(t testutil.TB, q *querier, expectedSeries []labels.Labels) {
   129  	t.Run("select", func(t testutil.TB) {
   130  		t.ResetTimer()
   131  
   132  		for i := 0; i < t.N(); i++ {
   133  			ss := q.Select(true, nil, &labels.Matcher{Value: "foo", Name: "bar", Type: labels.MatchEqual})
   134  			testutil.Ok(t, ss.Err())
   135  			testutil.Equals(t, 0, len(ss.Warnings()))
   136  
   137  			if t.IsBenchmark() {
   138  				var gotSeriesCount int
   139  				for ss.Next() {
   140  					s := ss.At()
   141  					testLset = s.Labels()
   142  					gotSeriesCount++
   143  
   144  					// This is when resource usage should actually start growing.
   145  					iter := s.Iterator(nil)
   146  					for iter.Next() != chunkenc.ValNone {
   147  						testT, testV = iter.At()
   148  					}
   149  					testutil.Ok(t, iter.Err())
   150  				}
   151  
   152  				testutil.Equals(t, len(expectedSeries), gotSeriesCount)
   153  				testutil.Ok(t, ss.Err())
   154  				return
   155  			}
   156  
   157  			// Check more carefully.
   158  			var gotSeries []labels.Labels
   159  			for ss.Next() {
   160  				s := ss.At()
   161  				gotSeries = append(gotSeries, s.Labels())
   162  
   163  				iter := s.Iterator(nil)
   164  				for iter.Next() != chunkenc.ValNone {
   165  					testT, testV = iter.At()
   166  				}
   167  				testutil.Ok(t, iter.Err())
   168  			}
   169  			testutil.Equals(t, expectedSeries, gotSeries)
   170  			testutil.Ok(t, ss.Err())
   171  		}
   172  	})
   173  }