github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/tsdb/index_test.go (about)

     1  package tsdb
     2  
     3  import (
     4  	"fmt"
     5  	"sort"
     6  	"testing"
     7  
     8  	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
     9  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    10  	"github.com/grafana/pyroscope/pkg/phlaredb/tsdb/shard"
    11  
    12  	"github.com/prometheus/common/model"
    13  	"github.com/prometheus/prometheus/model/labels"
    14  	"github.com/stretchr/testify/require"
    15  )
    16  
    17  func Test_GetShards(t *testing.T) {
    18  	for _, tt := range []struct {
    19  		total    uint32
    20  		shard    *shard.Annotation
    21  		expected []uint32
    22  	}{
    23  		// equal factors
    24  		{16, &shard.Annotation{Shard: 0, Of: 16}, []uint32{0}},
    25  		{16, &shard.Annotation{Shard: 4, Of: 16}, []uint32{4}},
    26  		{16, &shard.Annotation{Shard: 15, Of: 16}, []uint32{15}},
    27  
    28  		// idx factor a larger multiple of schema factor
    29  		{32, &shard.Annotation{Shard: 0, Of: 16}, []uint32{0, 16}},
    30  		{32, &shard.Annotation{Shard: 4, Of: 16}, []uint32{4, 20}},
    31  		{32, &shard.Annotation{Shard: 15, Of: 16}, []uint32{15, 31}},
    32  		{64, &shard.Annotation{Shard: 15, Of: 16}, []uint32{15, 31, 47, 63}},
    33  	} {
    34  		tt := tt
    35  		t.Run(tt.shard.String()+fmt.Sprintf("_total_%d", tt.total), func(t *testing.T) {
    36  			ii := NewWithShards(tt.total)
    37  			res := ii.getShards(tt.shard)
    38  			resInt := []uint32{}
    39  			for _, r := range res {
    40  				resInt = append(resInt, r.shard)
    41  			}
    42  			require.Equal(t, tt.expected, resInt)
    43  		})
    44  	}
    45  }
    46  
    47  func Test_ValidateShards(t *testing.T) {
    48  	ii := NewWithShards(32)
    49  	require.NoError(t, ii.validateShard(&shard.Annotation{Shard: 1, Of: 16}))
    50  }
    51  
    52  func TestDeleteAddLoopkup(t *testing.T) {
    53  	index := NewWithShards(DefaultIndexShards)
    54  	lbs := []*typesv1.LabelPair{
    55  		{Name: "__name__", Value: "foo"},
    56  		{Name: "foo", Value: "foo"},
    57  		{Name: "bar", Value: "bar"},
    58  		{Name: "buzz", Value: "buzz"},
    59  	}
    60  	sort.Sort(phlaremodel.Labels(lbs))
    61  
    62  	require.Equal(t, uint32(6), labelsSeriesIDHash(lbs)%32)
    63  	// make sure we consistent
    64  	require.Equal(t, uint32(6), labelsSeriesIDHash(lbs)%32)
    65  	index.Add(lbs, model.Fingerprint((phlaremodel.Labels(lbs).Hash())))
    66  	index.Delete(lbs, model.Fingerprint(phlaremodel.Labels(lbs).Hash()))
    67  	ids, err := index.Lookup([]*labels.Matcher{
    68  		labels.MustNewMatcher(labels.MatchEqual, "foo", "foo"),
    69  	}, nil)
    70  	require.NoError(t, err)
    71  	require.Len(t, ids, 0)
    72  }
    73  
    74  func Test_hash_mapping(t *testing.T) {
    75  	lbs := []*typesv1.LabelPair{
    76  		{Name: "compose_project", Value: "loki-boltdb-storage-s3"},
    77  		{Name: "compose_service", Value: "ingester-2"},
    78  		{Name: "container_name", Value: "loki-boltdb-storage-s3_ingester-2_1"},
    79  		{Name: "filename", Value: "/var/log/docker/790fef4c6a587c3b386fe85c07e03f3a1613f4929ca3abaa4880e14caadb5ad1/json.log"},
    80  		{Name: "host", Value: "docker-desktop"},
    81  		{Name: "source", Value: "stderr"},
    82  	}
    83  
    84  	for _, shardID := range []uint32{16, 32, 64, 128} {
    85  		t.Run(fmt.Sprintf("%d", shardID), func(t *testing.T) {
    86  			ii := NewWithShards(shardID)
    87  			ii.Add(lbs, 1)
    88  
    89  			res, err := ii.Lookup([]*labels.Matcher{{Type: labels.MatchEqual, Name: "compose_project", Value: "loki-boltdb-storage-s3"}}, &shard.Annotation{Shard: int(labelsSeriesIDHash(lbs) % 16), Of: 16})
    90  			require.NoError(t, err)
    91  			require.Len(t, res, 1)
    92  			require.Equal(t, model.Fingerprint(1), res[0])
    93  		})
    94  	}
    95  }
    96  
    97  func Test_NoMatcherLookup(t *testing.T) {
    98  	lbs := []*typesv1.LabelPair{
    99  		{Name: "foo", Value: "bar"},
   100  		{Name: "hi", Value: "hello"},
   101  	}
   102  	// with no shard param
   103  	ii := NewWithShards(16)
   104  	ii.Add(lbs, 1)
   105  	ids, err := ii.Lookup(nil, nil)
   106  	require.Nil(t, err)
   107  	require.Equal(t, model.Fingerprint(1), ids[0])
   108  
   109  	// with shard param
   110  	ii = NewWithShards(16)
   111  	ii.Add(lbs, 1)
   112  	ids, err = ii.Lookup(nil, &shard.Annotation{Shard: int(labelsSeriesIDHash(lbs) % 16), Of: 16})
   113  	require.Nil(t, err)
   114  	require.Equal(t, model.Fingerprint(1), ids[0])
   115  }
   116  
   117  func Test_ConsistentMapping(t *testing.T) {
   118  	a := NewWithShards(16)
   119  	b := NewWithShards(32)
   120  
   121  	for i := 0; i < 100; i++ {
   122  		lbs := []*typesv1.LabelPair{
   123  			{Name: "foo", Value: "bar"},
   124  			{Name: "hi", Value: fmt.Sprint(i)},
   125  		}
   126  		a.Add(lbs, model.Fingerprint(i))
   127  		b.Add(lbs, model.Fingerprint(i))
   128  	}
   129  
   130  	shardMax := 8
   131  	for i := 0; i < shardMax; i++ {
   132  		shard := &shard.Annotation{
   133  			Shard: i,
   134  			Of:    shardMax,
   135  		}
   136  
   137  		aIDs, err := a.Lookup([]*labels.Matcher{
   138  			labels.MustNewMatcher(labels.MatchEqual, "foo", "bar"),
   139  		}, shard)
   140  		require.Nil(t, err)
   141  
   142  		bIDs, err := b.Lookup([]*labels.Matcher{
   143  			labels.MustNewMatcher(labels.MatchEqual, "foo", "bar"),
   144  		}, shard)
   145  		require.Nil(t, err)
   146  
   147  		sorter := func(xs []model.Fingerprint) {
   148  			sort.Slice(xs, func(i, j int) bool {
   149  				return xs[i] < xs[j]
   150  			})
   151  		}
   152  		sorter(aIDs)
   153  		sorter(bIDs)
   154  
   155  		require.Equal(t, aIDs, bIDs, "incorrect shard mapping for shard %v", shard)
   156  	}
   157  }