github.com/grafana/pyroscope@v1.18.0/pkg/querybackend/query_series_labels.go (about)

     1  package querybackend
     2  
     3  import (
     4  	"errors"
     5  	"slices"
     6  	"sync"
     7  
     8  	"github.com/prometheus/prometheus/model/labels"
     9  	"github.com/prometheus/prometheus/storage"
    10  
    11  	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
    12  	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
    13  	"github.com/grafana/pyroscope/pkg/block"
    14  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    15  	"github.com/grafana/pyroscope/pkg/phlaredb"
    16  )
    17  
    18  func init() {
    19  	registerQueryType(
    20  		queryv1.QueryType_QUERY_SERIES_LABELS,
    21  		queryv1.ReportType_REPORT_SERIES_LABELS,
    22  		querySeriesLabels,
    23  		newSeriesLabelsAggregator,
    24  		false,
    25  		[]block.Section{block.SectionTSDB}...,
    26  	)
    27  }
    28  
    29  func querySeriesLabels(q *queryContext, query *queryv1.Query) (*queryv1.Report, error) {
    30  	series, err := getSeriesLabels(q.ds.Index(), q.req.matchers, query.SeriesLabels.LabelNames...)
    31  	if err != nil {
    32  		return nil, err
    33  	}
    34  	resp := &queryv1.Report{
    35  		SeriesLabels: &queryv1.SeriesLabelsReport{
    36  			Query:        query.SeriesLabels.CloneVT(),
    37  			SeriesLabels: series,
    38  		},
    39  	}
    40  	return resp, nil
    41  }
    42  
    43  type seriesLabelsAggregator struct {
    44  	init   sync.Once
    45  	query  *queryv1.SeriesLabelsQuery
    46  	series *phlaremodel.LabelMerger
    47  }
    48  
    49  func newSeriesLabelsAggregator(*queryv1.InvokeRequest) aggregator {
    50  	return new(seriesLabelsAggregator)
    51  }
    52  
    53  func (a *seriesLabelsAggregator) aggregate(report *queryv1.Report) error {
    54  	r := report.SeriesLabels
    55  	a.init.Do(func() {
    56  		a.query = r.Query.CloneVT()
    57  		a.series = phlaremodel.NewLabelMerger()
    58  	})
    59  	a.series.MergeSeries(r.SeriesLabels)
    60  	return nil
    61  }
    62  
    63  func (a *seriesLabelsAggregator) build() *queryv1.Report {
    64  	return &queryv1.Report{
    65  		SeriesLabels: &queryv1.SeriesLabelsReport{
    66  			Query:        a.query,
    67  			SeriesLabels: a.series.Labels(),
    68  		},
    69  	}
    70  }
    71  
    72  func getSeriesLabels(reader phlaredb.IndexReader, matchers []*labels.Matcher, by ...string) ([]*typesv1.Labels, error) {
    73  	names, err := reader.LabelNames()
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  	if len(by) > 0 {
    78  		names = slices.DeleteFunc(names, func(n string) bool {
    79  			for j := 0; j < len(by); j++ {
    80  				if by[j] == n {
    81  					return false
    82  				}
    83  			}
    84  			return true
    85  		})
    86  		if len(names) == 0 {
    87  			return nil, nil
    88  		}
    89  	}
    90  
    91  	postings, err := getPostings(reader, matchers...)
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  
    96  	visited := make(map[uint64]struct{})
    97  	sets := make([]*typesv1.Labels, 0, 32)
    98  	ls := make(phlaremodel.Labels, len(names))
    99  	for i := range names {
   100  		ls[i] = new(typesv1.LabelPair)
   101  	}
   102  
   103  	for postings.Next() {
   104  		var j int
   105  		var v string
   106  		for i := range names {
   107  			v, err = reader.LabelValueFor(postings.At(), names[i])
   108  			switch {
   109  			case err == nil:
   110  			case errors.Is(err, storage.ErrNotFound):
   111  			default:
   112  				return nil, err
   113  			}
   114  			if v != "" {
   115  				ls[j].Name = names[i]
   116  				ls[j].Value = v
   117  				j++
   118  			}
   119  		}
   120  		set := ls[:j]
   121  		h := set.Hash()
   122  		if _, ok := visited[h]; !ok {
   123  			visited[h] = struct{}{}
   124  			sets = append(sets, &typesv1.Labels{Labels: set.Clone()})
   125  		}
   126  	}
   127  
   128  	return sets, postings.Err()
   129  }