Move to dedicated hints types and add tests

Signed-off-by: Goutham Veeramachaneni <[email protected]>

Author: gouthamve

pkg/store/bucket.go

@@ -1059,7 +1059,7 @@ func chunksSize(chks []storepb.AggrChunk) (size int) {
 
 // LabelNames implements the storepb.StoreServer interface.
 func (s *BucketStore) LabelNames(ctx context.Context, req *storepb.LabelNamesRequest) (*storepb.LabelNamesResponse, error) {
-	resHints := &hintspb.SeriesResponseHints{}
+	resHints := &hintspb.LabelNamesResponseHints{}
 
 	g, gctx := errgroup.WithContext(ctx)
 
@@ -1113,7 +1113,7 @@ func (s *BucketStore) LabelNames(ctx context.Context, req *storepb.LabelNamesReq
 
 // LabelValues implements the storepb.StoreServer interface.
 func (s *BucketStore) LabelValues(ctx context.Context, req *storepb.LabelValuesRequest) (*storepb.LabelValuesResponse, error) {
-	resHints := &hintspb.SeriesResponseHints{}
+	resHints := &hintspb.LabelValuesResponseHints{}
 
 	g, gctx := errgroup.WithContext(ctx)
 

pkg/store/bucket_test.go

@@ -1897,3 +1897,190 @@ func createBlockWithOneSeriesWithStep(t testutil.TB, dir string, lbls labels.Lab
 
 	return createBlockFromHead(t, dir, h)
 }
+
+func TestLabelNamesAndValuesHints(t *testing.T) {
+	tb := testutil.NewTB(t)
+
+	tmpDir, err := ioutil.TempDir("", "test-labels-hints")
+	testutil.Ok(t, err)
+	defer func() { testutil.Ok(t, os.RemoveAll(tmpDir)) }()
+
+	bktDir := filepath.Join(tmpDir, "bkt")
+	bkt, err := filesystem.NewBucket(bktDir)
+	testutil.Ok(t, err)
+	defer func() { testutil.Ok(t, bkt.Close()) }()
+
+	var (
+		logger   = log.NewNopLogger()
+		instrBkt = objstore.WithNoopInstr(bkt)
+		random   = rand.New(rand.NewSource(120))
+	)
+
+	extLset := labels.Labels{{Name: "ext1", Value: "1"}}
+	// Inject the Thanos meta to each block in the storage.
+	thanosMeta := metadata.Thanos{
+		Labels:     extLset.Map(),
+		Downsample: metadata.ThanosDownsample{Resolution: 0},
+		Source:     metadata.TestSource,
+	}
+
+	// Create TSDB blocks.
+	head, seriesSet1 := storetestutil.CreateHeadWithSeries(t, 0, storetestutil.HeadGenOptions{
+		TSDBDir:          filepath.Join(tmpDir, "0"),
+		SamplesPerSeries: 1,
+		Series:           2,
+		PrependLabels:    extLset,
+		Random:           random,
+	})
+	block1 := createBlockFromHead(t, bktDir, head)
+	testutil.Ok(t, head.Close())
+	head2, seriesSet2 := storetestutil.CreateHeadWithSeries(t, 1, storetestutil.HeadGenOptions{
+		TSDBDir:          filepath.Join(tmpDir, "1"),
+		SamplesPerSeries: 1,
+		Series:           2,
+		PrependLabels:    extLset,
+		Random:           random,
+	})
+	block2 := createBlockFromHead(t, bktDir, head2)
+	testutil.Ok(t, head2.Close())
+
+	for _, blockID := range []ulid.ULID{block1, block2} {
+		_, err := metadata.InjectThanos(logger, filepath.Join(bktDir, blockID.String()), thanosMeta, nil)
+		testutil.Ok(t, err)
+	}
+
+	// Instance a real bucket store we'll use to query back the series.
+	fetcher, err := block.NewMetaFetcher(logger, 10, instrBkt, tmpDir, nil, nil, nil)
+	testutil.Ok(tb, err)
+
+	indexCache, err := storecache.NewInMemoryIndexCacheWithConfig(logger, nil, storecache.InMemoryIndexCacheConfig{})
+	testutil.Ok(tb, err)
+
+	store, err := NewBucketStore(
+		logger,
+		nil,
+		instrBkt,
+		fetcher,
+		tmpDir,
+		indexCache,
+		nil,
+		1000000,
+		NewChunksLimiterFactory(10000/MaxSamplesPerChunk),
+		false,
+		10,
+		nil,
+		false,
+		true,
+		DefaultPostingOffsetInMemorySampling,
+		true,
+	)
+	testutil.Ok(tb, err)
+	testutil.Ok(tb, store.SyncBlocks(context.Background()))
+
+	type labelNamesValuesCase struct {
+		name string
+
+		labelNamesReq      *storepb.LabelNamesRequest
+		expectedNames      []string
+		expectedNamesHints hintspb.LabelNamesResponseHints
+
+		labelValuesReq      *storepb.LabelValuesRequest
+		expectedValues      []string
+		expectedValuesHints hintspb.LabelValuesResponseHints
+	}
+
+	testCases := []labelNamesValuesCase{
+		{
+			name: "querying a range containing 1 block should return 1 block in the labels hints",
+
+			labelNamesReq: &storepb.LabelNamesRequest{
+				Start: 0,
+				End:   1,
+			},
+			expectedNames: labelNamesFromSeriesSet(seriesSet1),
+			expectedNamesHints: hintspb.LabelNamesResponseHints{
+				QueriedBlocks: []hintspb.Block{
+					{Id: block1.String()},
+				},
+			},
+
+			labelValuesReq: &storepb.LabelValuesRequest{
+				Label: "__name__",
+				Start: 0,
+				End:   1,
+			},
+			expectedValues: []string{},
+			expectedValuesHints: hintspb.LabelValuesResponseHints{
+				QueriedBlocks: []hintspb.Block{
+					{Id: block1.String()},
+				},
+			},
+		},
+		{
+			name: "querying a range containing multiple blocks should return multiple blocks in the response hints",
+
+			labelNamesReq: &storepb.LabelNamesRequest{
+				Start: 0,
+				End:   3,
+			},
+			expectedNames: labelNamesFromSeriesSet(
+				append(append([]*storepb.Series{}, seriesSet1...), seriesSet2...),
+			),
+			expectedNamesHints: hintspb.LabelNamesResponseHints{
+				QueriedBlocks: []hintspb.Block{
+					{Id: block1.String()},
+					{Id: block2.String()},
+				},
+			},
+
+			labelValuesReq: &storepb.LabelValuesRequest{
+				Label: "__name__",
+				Start: 0,
+				End:   3,
+			},
+			expectedValues: []string{},
+			expectedValuesHints: hintspb.LabelValuesResponseHints{
+				QueriedBlocks: []hintspb.Block{
+					{Id: block1.String()},
+					{Id: block2.String()},
+				},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			namesResp, err := store.LabelNames(context.Background(), tc.labelNamesReq)
+			testutil.Ok(t, err)
+			testutil.Equals(t, tc.expectedNames, namesResp.Names)
+			var namesHints hintspb.LabelNamesResponseHints
+			testutil.Ok(t, types.UnmarshalAny(namesResp.Hints, &namesHints))
+			testutil.Equals(t, tc.expectedNamesHints, namesHints)
+
+			valuesResp, err := store.LabelValues(context.Background(), tc.labelValuesReq)
+			testutil.Ok(t, err)
+			testutil.Equals(t, tc.expectedValues, valuesResp.Values)
+			var valuesHints hintspb.LabelValuesResponseHints
+			testutil.Ok(t, types.UnmarshalAny(valuesResp.Hints, &valuesHints))
+			testutil.Equals(t, tc.expectedValuesHints, valuesHints)
+		})
+	}
+}
+
+func labelNamesFromSeriesSet(series []*storepb.Series) []string {
+	labelsMap := map[string]struct{}{}
+
+	for _, s := range series {
+		for _, label := range s.Labels {
+			labelsMap[label.Name] = struct{}{}
+		}
+	}
+
+	labels := make([]string, 0, len(labelsMap))
+	for k := range labelsMap {
+		labels = append(labels, k)
+	}
+
+	sort.Strings(labels)
+	return labels
+}

pkg/store/hintspb/custom.go

@@ -10,3 +10,15 @@ func (m *SeriesResponseHints) AddQueriedBlock(id ulid.ULID) {
 		Id: id.String(),
 	})
 }
+
+func (m *LabelNamesResponseHints) AddQueriedBlock(id ulid.ULID) {
+	m.QueriedBlocks = append(m.QueriedBlocks, Block{
+		Id: id.String(),
+	})
+}
+
+func (m *LabelValuesResponseHints) AddQueriedBlock(id ulid.ULID) {
+	m.QueriedBlocks = append(m.QueriedBlocks, Block{
+		Id: id.String(),
+	})
+}