github.com/grafana/pyroscope@v1.18.0/pkg/querybackend/query_tree.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  
    11  	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
    12  	"github.com/grafana/pyroscope/pkg/block"
    13  	"github.com/grafana/pyroscope/pkg/model"
    14  	parquetquery "github.com/grafana/pyroscope/pkg/phlaredb/query"
    15  	v1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
    16  	"github.com/grafana/pyroscope/pkg/phlaredb/symdb"
    17  )
    18  
    19  func init() {
    20  	registerQueryType(
    21  		queryv1.QueryType_QUERY_TREE,
    22  		queryv1.ReportType_REPORT_TREE,
    23  		queryTree,
    24  		newTreeAggregator,
    25  		false,
    26  		[]block.Section{
    27  			block.SectionTSDB,
    28  			block.SectionProfiles,
    29  			block.SectionSymbols,
    30  		}...,
    31  	)
    32  }
    33  
    34  func queryTree(q *queryContext, query *queryv1.Query) (*queryv1.Report, error) {
    35  	span := opentracing.SpanFromContext(q.ctx)
    36  
    37  	var profileOpts []profileIteratorOption
    38  	if len(query.Tree.ProfileIdSelector) > 0 {
    39  		opt, err := withProfileIDSelector(query.Tree.ProfileIdSelector...)
    40  		if err != nil {
    41  			return nil, err
    42  		}
    43  		profileOpts = append(profileOpts, opt)
    44  		span.SetTag("profile_id_selector.count", len(query.Tree.ProfileIdSelector))
    45  		if len(query.Tree.ProfileIdSelector) <= maxProfileIDsToLog {
    46  			span.LogFields(otlog.String("profile_ids", strings.Join(query.Tree.ProfileIdSelector, ",")))
    47  		}
    48  	}
    49  
    50  	entries, err := profileEntryIterator(q, profileOpts...)
    51  	if err != nil {
    52  		return nil, err
    53  	}
    54  	defer runutil.CloseWithErrCapture(&err, entries, "failed to close profile entry iterator")
    55  
    56  	spanSelector, err := model.NewSpanSelector(query.Tree.SpanSelector)
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  
    61  	var columns v1.SampleColumns
    62  	if err = columns.Resolve(q.ds.Profiles().Schema()); err != nil {
    63  		return nil, err
    64  	}
    65  
    66  	indices := []int{
    67  		columns.StacktraceID.ColumnIndex,
    68  		columns.Value.ColumnIndex,
    69  	}
    70  	if len(spanSelector) > 0 {
    71  		indices = append(indices, columns.SpanID.ColumnIndex)
    72  	}
    73  
    74  	resolverOptions := []symdb.ResolverOption{
    75  		symdb.WithResolverMaxNodes(query.Tree.MaxNodes),
    76  	}
    77  	if query.Tree.StackTraceSelector != nil {
    78  		resolverOptions = append(resolverOptions, symdb.WithResolverStackTraceSelector(query.Tree.StackTraceSelector))
    79  	}
    80  
    81  	profiles := parquetquery.NewRepeatedRowIterator(q.ctx, entries, q.ds.Profiles().RowGroups(), indices...)
    82  	defer runutil.CloseWithErrCapture(&err, profiles, "failed to close profile stream")
    83  
    84  	resolver := symdb.NewResolver(q.ctx, q.ds.Symbols(), resolverOptions...)
    85  	defer resolver.Release()
    86  
    87  	if len(spanSelector) > 0 {
    88  		for profiles.Next() {
    89  			p := profiles.At()
    90  			resolver.AddSamplesWithSpanSelectorFromParquetRow(
    91  				p.Row.Partition,
    92  				p.Values[0],
    93  				p.Values[1],
    94  				p.Values[2],
    95  				spanSelector,
    96  			)
    97  		}
    98  	} else {
    99  		for profiles.Next() {
   100  			p := profiles.At()
   101  			resolver.AddSamplesFromParquetRow(p.Row.Partition, p.Values[0], p.Values[1])
   102  		}
   103  	}
   104  
   105  	if err = profiles.Err(); err != nil {
   106  		return nil, err
   107  	}
   108  
   109  	tree, err := resolver.Tree()
   110  	if err != nil {
   111  		return nil, err
   112  	}
   113  
   114  	resp := &queryv1.Report{
   115  		Tree: &queryv1.TreeReport{
   116  			Query: query.Tree.CloneVT(),
   117  			Tree:  tree.Bytes(query.Tree.GetMaxNodes()),
   118  		},
   119  	}
   120  	return resp, nil
   121  }
   122  
   123  type treeAggregator struct {
   124  	init  sync.Once
   125  	query *queryv1.TreeQuery
   126  	tree  *model.TreeMerger
   127  }
   128  
   129  func newTreeAggregator(*queryv1.InvokeRequest) aggregator { return new(treeAggregator) }
   130  
   131  func (a *treeAggregator) aggregate(report *queryv1.Report) error {
   132  	r := report.Tree
   133  	a.init.Do(func() {
   134  		a.tree = model.NewTreeMerger()
   135  		a.query = r.Query.CloneVT()
   136  	})
   137  	return a.tree.MergeTreeBytes(r.Tree)
   138  }
   139  
   140  func (a *treeAggregator) build() *queryv1.Report {
   141  	return &queryv1.Report{
   142  		Tree: &queryv1.TreeReport{
   143  			Query: a.query,
   144  			Tree:  a.tree.Tree().Bytes(a.query.GetMaxNodes()),
   145  		},
   146  	}
   147  }