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

     1  package querybackend
     2  
     3  import (
     4  	"strings"
     5  	"sync"
     6  
     7  	"github.com/grafana/dskit/runutil"
     8  	"github.com/opentracing/opentracing-go"
     9  	otlog "github.com/opentracing/opentracing-go/log"
    10  	"github.com/prometheus/prometheus/model/labels"
    11  
    12  	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
    13  	"github.com/grafana/pyroscope/pkg/block"
    14  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    15  	parquetquery "github.com/grafana/pyroscope/pkg/phlaredb/query"
    16  	v1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
    17  	"github.com/grafana/pyroscope/pkg/phlaredb/symdb"
    18  	"github.com/grafana/pyroscope/pkg/pprof"
    19  )
    20  
    21  func init() {
    22  	registerQueryType(
    23  		queryv1.QueryType_QUERY_PPROF,
    24  		queryv1.ReportType_REPORT_PPROF,
    25  		queryPprof,
    26  		newPprofAggregator,
    27  		false,
    28  		[]block.Section{
    29  			block.SectionTSDB,
    30  			block.SectionProfiles,
    31  			block.SectionSymbols,
    32  		}...,
    33  	)
    34  }
    35  
    36  func queryPprof(q *queryContext, query *queryv1.Query) (*queryv1.Report, error) {
    37  	span := opentracing.SpanFromContext(q.ctx)
    38  
    39  	var profileOpts []profileIteratorOption
    40  	if len(query.Pprof.ProfileIdSelector) > 0 {
    41  		opt, err := withProfileIDSelector(query.Pprof.ProfileIdSelector...)
    42  		if err != nil {
    43  			return nil, err
    44  		}
    45  		profileOpts = append(profileOpts, opt)
    46  		span.SetTag("profile_id_selector.count", len(query.Pprof.ProfileIdSelector))
    47  		if len(query.Pprof.ProfileIdSelector) <= maxProfileIDsToLog {
    48  			span.LogFields(otlog.String("profile_ids", strings.Join(query.Pprof.ProfileIdSelector, ",")))
    49  		}
    50  	}
    51  
    52  	entries, err := profileEntryIterator(q, profileOpts...)
    53  	if err != nil {
    54  		return nil, err
    55  	}
    56  	defer runutil.CloseWithErrCapture(&err, entries, "failed to close profile entry iterator")
    57  
    58  	var columns v1.SampleColumns
    59  	if err = columns.Resolve(q.ds.Profiles().Schema()); err != nil {
    60  		return nil, err
    61  	}
    62  
    63  	profiles := parquetquery.NewRepeatedRowIterator(q.ctx, entries, q.ds.Profiles().RowGroups(),
    64  		columns.StacktraceID.ColumnIndex,
    65  		columns.Value.ColumnIndex)
    66  	defer runutil.CloseWithErrCapture(&err, profiles, "failed to close profile stream")
    67  
    68  	resolverOptions := make([]symdb.ResolverOption, 0)
    69  	resolverOptions = append(resolverOptions, symdb.WithResolverMaxNodes(query.Pprof.MaxNodes))
    70  	if query.Pprof.StackTraceSelector != nil {
    71  		resolverOptions = append(
    72  			resolverOptions,
    73  			symdb.WithResolverStackTraceSelector(query.Pprof.StackTraceSelector),
    74  			symdb.WithResolverSanitizeOnMerge(q.req.src.Options.SanitizeOnMerge))
    75  	}
    76  
    77  	resolver := symdb.NewResolver(q.ctx, q.ds.Symbols(), resolverOptions...)
    78  	defer resolver.Release()
    79  
    80  	for profiles.Next() {
    81  		p := profiles.At()
    82  		resolver.AddSamplesFromParquetRow(p.Row.Partition, p.Values[0], p.Values[1])
    83  	}
    84  	if err = profiles.Err(); err != nil {
    85  		return nil, err
    86  	}
    87  
    88  	profile, err := resolver.Pprof()
    89  	if err != nil {
    90  		return nil, err
    91  	}
    92  
    93  	for _, m := range q.req.matchers {
    94  		if m.Name == phlaremodel.LabelNameProfileType && m.Type == labels.MatchEqual {
    95  			if t, err := phlaremodel.ParseProfileTypeSelector(m.Value); err == nil {
    96  				pprof.SetProfileMetadata(profile, t, q.req.endTime, 0)
    97  				break
    98  			}
    99  		}
   100  	}
   101  
   102  	resp := &queryv1.Report{
   103  		Pprof: &queryv1.PprofReport{
   104  			Query: query.Pprof.CloneVT(),
   105  			Pprof: pprof.MustMarshal(profile, true),
   106  		},
   107  	}
   108  
   109  	return resp, nil
   110  }
   111  
   112  type pprofAggregator struct {
   113  	init     sync.Once
   114  	query    *queryv1.PprofQuery
   115  	profile  pprof.ProfileMerge
   116  	sanitize bool
   117  }
   118  
   119  func newPprofAggregator(req *queryv1.InvokeRequest) aggregator {
   120  	if req.Options != nil {
   121  		return &pprofAggregator{
   122  			sanitize: req.Options.SanitizeOnMerge,
   123  		}
   124  	}
   125  	return &pprofAggregator{}
   126  }
   127  
   128  func (a *pprofAggregator) aggregate(report *queryv1.Report) error {
   129  	r := report.Pprof
   130  	a.init.Do(func() {
   131  		a.query = r.Query.CloneVT()
   132  	})
   133  	return a.profile.MergeBytes(r.Pprof, a.sanitize)
   134  }
   135  
   136  func (a *pprofAggregator) build() *queryv1.Report {
   137  	return &queryv1.Report{
   138  		Pprof: &queryv1.PprofReport{
   139  			Query: a.query,
   140  			Pprof: pprof.MustMarshal(a.profile.Profile(), true),
   141  		},
   142  	}
   143  }