github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/tools/tsdb/tsdb-map/main_test.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"math"
     7  	"math/rand"
     8  	"os"
     9  	"testing"
    10  
    11  	"github.com/prometheus/prometheus/model/labels"
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"github.com/grafana/loki/pkg/storage/stores/tsdb"
    15  	"github.com/grafana/loki/pkg/storage/stores/tsdb/index"
    16  )
    17  
    18  func TestExtractChecksum(t *testing.T) {
    19  	x := rand.Uint32()
    20  	s := fmt.Sprintf("a/b/c:d:e:%x", x)
    21  	require.Equal(t, x, extractChecksumFromChunkID([]byte(s)))
    22  }
    23  
    24  type testCase struct {
    25  	name     string
    26  	matchers []*labels.Matcher
    27  }
    28  
    29  var cases = []testCase{
    30  	{
    31  		name:     "match ns",
    32  		matchers: []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "namespace", "loki-ops")},
    33  	},
    34  	{
    35  		name:     "match job regexp",
    36  		matchers: []*labels.Matcher{labels.MustNewMatcher(labels.MatchRegexp, "job", "loki.*/distributor")},
    37  	},
    38  	{
    39  		name: "match regexp and equals",
    40  		matchers: []*labels.Matcher{
    41  			labels.MustNewMatcher(labels.MatchRegexp, "job", "loki.*/distributor"),
    42  			labels.MustNewMatcher(labels.MatchEqual, "cluster", "prod-us-central-0"),
    43  		},
    44  	},
    45  	{
    46  		name: "not equals cluster",
    47  		matchers: []*labels.Matcher{
    48  			labels.MustNewMatcher(labels.MatchNotEqual, "cluster", "prod-us-central-0"),
    49  		},
    50  	},
    51  	{
    52  		name: "inverse regexp",
    53  		matchers: []*labels.Matcher{
    54  			labels.MustNewMatcher(labels.MatchNotRegexp, "cluster", "prod.*"),
    55  		},
    56  	},
    57  }
    58  
    59  // Only iterates through the low level IndexReader using `PostingsForMatchers`
    60  // Requires LOKI_TSDB_PATH to be set or this will short-circuit
    61  func BenchmarkQuery_PostingsForMatchers(b *testing.B) {
    62  	for _, bm := range cases {
    63  		indexPath := os.Getenv("LOKI_TSDB_PATH")
    64  		if indexPath == "" {
    65  			return
    66  		}
    67  
    68  		reader, err := index.NewFileReader(indexPath)
    69  		if err != nil {
    70  			panic(err)
    71  		}
    72  		b.Run(bm.name, func(b *testing.B) {
    73  			for i := 0; i < b.N; i++ {
    74  				p, _ := tsdb.PostingsForMatchers(reader, nil, bm.matchers...)
    75  
    76  				for p.Next() {
    77  				}
    78  			}
    79  		})
    80  	}
    81  }
    82  
    83  // Uses the higher level loki index interface, resolving chunk refs.
    84  // Requires LOKI_TSDB_PATH to be set or this will short-circuit
    85  func BenchmarkQuery_GetChunkRefs(b *testing.B) {
    86  	for _, bm := range cases {
    87  		indexPath := os.Getenv("LOKI_TSDB_PATH")
    88  		if indexPath == "" {
    89  			return
    90  		}
    91  
    92  		reader, err := index.NewFileReader(indexPath)
    93  		if err != nil {
    94  			panic(err)
    95  		}
    96  		idx := tsdb.NewTSDBIndex(reader)
    97  		b.Run(bm.name, func(b *testing.B) {
    98  			refs := tsdb.ChunkRefsPool.Get()
    99  			for i := 0; i < b.N; i++ {
   100  				var err error
   101  				refs, err = idx.GetChunkRefs(context.Background(), "fake", 0, math.MaxInt64, refs, nil, bm.matchers...)
   102  				if err != nil {
   103  					panic(err)
   104  				}
   105  			}
   106  		})
   107  	}
   108  }
   109  
   110  func BenchmarkQuery_GetChunkRefsSharded(b *testing.B) {
   111  	for _, bm := range cases {
   112  		indexPath := os.Getenv("LOKI_TSDB_PATH")
   113  		if indexPath == "" {
   114  			return
   115  		}
   116  
   117  		reader, err := index.NewFileReader(indexPath)
   118  		if err != nil {
   119  			panic(err)
   120  		}
   121  		idx := tsdb.NewTSDBIndex(reader)
   122  		shardFactor := 16
   123  
   124  		b.Run(bm.name, func(b *testing.B) {
   125  			refs := tsdb.ChunkRefsPool.Get()
   126  			for i := 0; i < b.N; i++ {
   127  				for j := 0; j < shardFactor; j++ {
   128  					shard := index.ShardAnnotation{
   129  						Shard: uint32(j),
   130  						Of:    uint32(shardFactor),
   131  					}
   132  					var err error
   133  
   134  					refs, err = idx.GetChunkRefs(context.Background(), "fake", 0, math.MaxInt64, refs, &shard, bm.matchers...)
   135  					if err != nil {
   136  						panic(err)
   137  					}
   138  				}
   139  			}
   140  		})
   141  	}
   142  }