perf(tsdb): don't fetch labels in TSDBIndex.GetChunkRefs (#20970)

Signed-off-by: Bryan Boreham <bjboreham@gmail.com>

Author: bboreham

pkg/storage/stores/shipper/indexshipper/tsdb/head_read.go

@@ -111,14 +111,17 @@ func (h *headIndexReader) Postings(name string, fpFilter index.FingerprintFilter
 }
 
 // Series returns the series for the given reference.
+// lbls can be nil, to indicate that just the chunks are needed.
 func (h *headIndexReader) Series(ref storage.SeriesRef, from int64, through int64, lbls *labels.Labels, chks *[]index.ChunkMeta) (uint64, error) {
 	s := h.head.series.getByID(uint64(ref))
 
 	if s == nil {
 		h.head.metrics.seriesNotFound.Inc()
 		return 0, storage.ErrNotFound
 	}
-	lbls.CopyFrom(s.ls)
+	if lbls != nil {
+		lbls.CopyFrom(s.ls)
+	}
 
 	queryBounds := newBounds(model.Time(from), model.Time(through))
 

pkg/storage/stores/shipper/indexshipper/tsdb/index/index.go

@@ -2135,14 +2135,12 @@ func buildChunkSamples(d encoding.Decbuf, numChunks int, info *chunkSamples) err
 	return d.Err()
 }
 
-func (dec *Decoder) prepSeries(b []byte, lbls *labels.Labels, chks *[]ChunkMeta) (*encoding.Decbuf, uint64, error) {
+// prepSeries returns series labels for a series, only returning selected `by` label names.
+// If `by` is nil, it returns all labels for the series.
+func (dec *Decoder) prepSeries(b []byte, lbls *labels.Labels, by map[string]struct{}) (*encoding.Decbuf, uint64, error) {
 	builder := labelpool.Get()
 	defer labelpool.Put(builder)
 
-	if chks != nil {
-		*chks = (*chks)[:0]
-	}
-
 	d := encoding.DecWrap(tsdb_enc.Decbuf{B: b})
 
 	fprint := d.Be64()
@@ -2160,6 +2158,13 @@ func (dec *Decoder) prepSeries(b []byte, lbls *labels.Labels, chks *[]ChunkMeta)
 		if err != nil {
 			return nil, 0, errors.Wrap(err, "lookup label name")
 		}
+
+		if by != nil {
+			if _, ok := by[ln]; !ok {
+				continue
+			}
+		}
+
 		lv, err := dec.LookupSymbol(lvo)
 		if err != nil {
 			return nil, 0, errors.Wrap(err, "lookup label value")
@@ -2171,62 +2176,30 @@ func (dec *Decoder) prepSeries(b []byte, lbls *labels.Labels, chks *[]ChunkMeta)
 	// Commit built labels.
 	builder.Sort()
 	*lbls = builder.Labels()
-
 	return &d, fprint, nil
 }
 
-// prepSeriesBy returns series labels and chunks for a series and only returning selected `by` label names.
-// If `by` is empty, it returns all labels for the series.
-func (dec *Decoder) prepSeriesBy(b []byte, lbls *labels.Labels, chks *[]ChunkMeta, by map[string]struct{}) (*encoding.Decbuf, uint64, error) {
-	if by == nil {
-		return dec.prepSeries(b, lbls, chks)
-	}
-
-	builder := labelpool.Get()
-	defer labelpool.Put(builder)
-
-	if chks != nil {
-		*chks = (*chks)[:0]
-	}
-
+// skipSeriesLabels reads past the label section in buffer b, ready to read chunks after that.
+func (dec *Decoder) skipSeriesLabels(b []byte) (*encoding.Decbuf, uint64, error) {
 	d := encoding.DecWrap(tsdb_enc.Decbuf{B: b})
 
 	fprint := d.Be64()
 	k := d.Uvarint()
 
-	for i := 0; i < k; i++ {
-		lno := uint32(d.Uvarint())
-		lvo := uint32(d.Uvarint())
+	for range k {
+		_ = d.Uvarint()
+		_ = d.Uvarint()
 
 		if d.Err() != nil {
 			return nil, 0, errors.Wrap(d.Err(), "read series label offsets")
 		}
-		// todo(cyriltovena): we could cache this by user requests spanning multiple prepSeries calls.
-		ln, err := dec.LookupSymbol(lno)
-		if err != nil {
-			return nil, 0, errors.Wrap(err, "lookup label name")
-		}
-		if _, ok := by[ln]; !ok {
-			continue
-		}
-
-		lv, err := dec.LookupSymbol(lvo)
-		if err != nil {
-			return nil, 0, errors.Wrap(err, "lookup label value")
-		}
-
-		builder.Add(ln, lv)
 	}
 
-	// Commit built labels.
-	builder.Sort()
-	*lbls = builder.Labels()
-
 	return &d, fprint, nil
 }
 
 func (dec *Decoder) ChunkStats(version int, b []byte, seriesRef storage.SeriesRef, from, through int64, lbls *labels.Labels, by map[string]struct{}) (uint64, ChunkStats, error) {
-	d, fp, err := dec.prepSeriesBy(b, lbls, nil, by)
+	d, fp, err := dec.prepSeries(b, lbls, by)
 	if err != nil {
 		return 0, ChunkStats{}, err
 	}
@@ -2368,15 +2341,25 @@ func (dec *Decoder) readChunkStatsPriorV3(d *encoding.Decbuf, seriesRef storage.
 	return res, nil
 }
 
-// Series decodes a series entry from the given byte slice into lset and chks.
-func (dec *Decoder) Series(version int, b []byte, seriesRef storage.SeriesRef, from int64, through int64, lbls *labels.Labels, chks *[]ChunkMeta) (uint64, error) {
-	d, fprint, err := dec.prepSeries(b, lbls, chks)
+// Series decodes a series entry from the given byte slice into lbls and chks.
+// lbls can be nil, indicating the caller only wants the chunks.
+func (dec *Decoder) Series(version int, b []byte, seriesRef storage.SeriesRef, from int64, through int64, lbls *labels.Labels, chks *[]ChunkMeta) (fprint uint64, err error) {
+	var d *encoding.Decbuf
+	if lbls == nil {
+		d, fprint, err = dec.skipSeriesLabels(b)
+	} else {
+		d, fprint, err = dec.prepSeries(b, lbls, nil)
+	}
 	if err != nil {
 		return 0, err
 	}
 
+	*chks = (*chks)[:0]
 	// read chunks based on fmt
 	if err := dec.readChunks(version, d, seriesRef, from, through, chks); err != nil {
+		if lbls == nil {
+			return 0, errors.Wrapf(err, "series footprint %x", fprint)
+		}
 		return 0, errors.Wrapf(err, "series %s", lbls.String())
 	}
 	return fprint, nil

pkg/storage/stores/shipper/indexshipper/tsdb/single_file_index.go

@@ -162,17 +162,20 @@ func (i *TSDBIndex) SetChunkFilterer(chunkFilter chunk.RequestChunkFilterer) {
 // Accepts a userID argument in order to implement `Index` interface, but since this is a single tenant index,
 // it is ignored (it's enforced elsewhere in index selection)
 func (i *TSDBIndex) ForSeries(ctx context.Context, _ string, fpFilter index.FingerprintFilter, from model.Time, through model.Time, fn func(labels.Labels, model.Fingerprint, []index.ChunkMeta) (stop bool), matchers ...*labels.Matcher) error {
+	var filterer chunk.Filterer
+	if i.chunkFilter != nil {
+		filterer = i.chunkFilter.ForRequest(ctx)
+	}
+	return i.forSeriesAndLabels(ctx, fpFilter, filterer, from, through, fn, matchers...)
+}
+
+func (i *TSDBIndex) forSeriesAndLabels(ctx context.Context, fpFilter index.FingerprintFilter, filterer chunk.Filterer, from model.Time, through model.Time, fn func(labels.Labels, model.Fingerprint, []index.ChunkMeta) (stop bool), matchers ...*labels.Matcher) error {
 	// TODO(owen-d): use pool
 
 	var ls labels.Labels
 	chks := ChunkMetasPool.Get()
 	defer func() { ChunkMetasPool.Put(chks) }()
 
-	var filterer chunk.Filterer
-	if i.chunkFilter != nil {
-		filterer = i.chunkFilter.ForRequest(ctx)
-	}
-
 	return i.forPostings(ctx, fpFilter, from, through, matchers, func(p index.Postings) error {
 		for p.Next() {
 			hash, err := i.reader.Series(p.At(), int64(from), int64(through), &ls, &chks)
@@ -198,6 +201,31 @@ func (i *TSDBIndex) ForSeries(ctx context.Context, _ string, fpFilter index.Fing
 
 }
 
+// Same as ForSeries, but the callback fn does not take a Labels parameter.
+func (i *TSDBIndex) forSeriesNoLabels(ctx context.Context, fpFilter index.FingerprintFilter, from model.Time, through model.Time, fn func(model.Fingerprint, []index.ChunkMeta) (stop bool), matchers ...*labels.Matcher) error {
+	chks := ChunkMetasPool.Get()
+	defer ChunkMetasPool.Put(chks)
+
+	return i.forPostings(ctx, fpFilter, from, through, matchers, func(p index.Postings) error {
+		for p.Next() {
+			hash, err := i.reader.Series(p.At(), int64(from), int64(through), nil, &chks)
+			if err != nil {
+				return err
+			}
+
+			// skip series that belong to different shards
+			if fpFilter != nil && !fpFilter.Match(model.Fingerprint(hash)) {
+				continue
+			}
+
+			if stop := fn(model.Fingerprint(hash), chks); stop {
+				break
+			}
+		}
+		return p.Err()
+	})
+}
+
 func (i *TSDBIndex) forPostings(
 	_ context.Context,
 	fpFilter index.FingerprintFilter,
@@ -218,7 +246,7 @@ func (i *TSDBIndex) GetChunkRefs(ctx context.Context, userID string, from, throu
 	}
 	res = res[:0]
 
-	if err := i.ForSeries(ctx, "", fpFilter, from, through, func(_ labels.Labels, fp model.Fingerprint, chks []index.ChunkMeta) (stop bool) {
+	addChunksToResult := func(fp model.Fingerprint, chks []index.ChunkMeta) (stop bool) {
 		for _, chk := range chks {
 			res = append(res, logproto.ChunkRefWithSizingInfo{
 				ChunkRef: logproto.ChunkRef{
@@ -233,11 +261,22 @@ func (i *TSDBIndex) GetChunkRefs(ctx context.Context, userID string, from, throu
 			})
 		}
 		return false
-	}, matchers...); err != nil {
-		return nil, err
 	}
 
-	return res, nil
+	var filterer chunk.Filterer
+	if i.chunkFilter != nil {
+		filterer = i.chunkFilter.ForRequest(ctx)
+	}
+	var err error
+	if filterer != nil {
+		// We need to fetch labels to pass to the filterer, even though we don't look at them in the callback.
+		err = i.forSeriesAndLabels(ctx, fpFilter, filterer, from, through, func(_ labels.Labels, fp model.Fingerprint, chks []index.ChunkMeta) (stop bool) {
+			return addChunksToResult(fp, chks)
+		}, matchers...)
+	} else {
+		err = i.forSeriesNoLabels(ctx, fpFilter, from, through, addChunksToResult, matchers...)
+	}
+	return res, err
 }
 
 func (i *TSDBIndex) Series(ctx context.Context, _ string, from, through model.Time, res []Series, fpFilter index.FingerprintFilter, matchers ...*labels.Matcher) ([]Series, error) {