github.com/grafana/pyroscope@v1.18.0/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces.go (about)

     1  package queryfrontend
     2  
     3  import (
     4  	"context"
     5  
     6  	"connectrpc.com/connect"
     7  	"github.com/grafana/dskit/tenant"
     8  
     9  	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
    10  	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
    11  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    12  	"github.com/grafana/pyroscope/pkg/validation"
    13  )
    14  
    15  func (q *QueryFrontend) SelectMergeStacktraces(
    16  	ctx context.Context,
    17  	c *connect.Request[querierv1.SelectMergeStacktracesRequest],
    18  ) (*connect.Response[querierv1.SelectMergeStacktracesResponse], error) {
    19  	b, err := q.selectMergeStacktracesTree(ctx, c)
    20  	if err != nil {
    21  		return nil, err
    22  	}
    23  	var resp querierv1.SelectMergeStacktracesResponse
    24  	switch c.Msg.Format {
    25  	case querierv1.ProfileFormat_PROFILE_FORMAT_TREE:
    26  		resp.Tree = b
    27  	default:
    28  		t, err := phlaremodel.UnmarshalTree(b)
    29  		if err != nil {
    30  			return nil, err
    31  		}
    32  		resp.Flamegraph = phlaremodel.NewFlameGraph(t, c.Msg.GetMaxNodes())
    33  	}
    34  	return connect.NewResponse(&resp), nil
    35  }
    36  
    37  func (q *QueryFrontend) selectMergeStacktracesTree(
    38  	ctx context.Context,
    39  	c *connect.Request[querierv1.SelectMergeStacktracesRequest],
    40  ) (tree []byte, err error) {
    41  	tenantIDs, err := tenant.TenantIDs(ctx)
    42  	if err != nil {
    43  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    44  	}
    45  	empty, err := validation.SanitizeTimeRange(q.limits, tenantIDs, &c.Msg.Start, &c.Msg.End)
    46  	if err != nil {
    47  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    48  	}
    49  	if empty {
    50  		return nil, nil
    51  	}
    52  
    53  	maxNodes, err := validation.ValidateMaxNodes(q.limits, tenantIDs, c.Msg.GetMaxNodes())
    54  	if err != nil {
    55  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    56  	}
    57  
    58  	_, err = phlaremodel.ParseProfileTypeSelector(c.Msg.ProfileTypeID)
    59  	if err != nil {
    60  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    61  	}
    62  
    63  	labelSelector, err := buildLabelSelectorWithProfileType(c.Msg.LabelSelector, c.Msg.ProfileTypeID)
    64  	if err != nil {
    65  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    66  	}
    67  	report, err := q.querySingle(ctx, &queryv1.QueryRequest{
    68  		StartTime:     c.Msg.Start,
    69  		EndTime:       c.Msg.End,
    70  		LabelSelector: labelSelector,
    71  		Query: []*queryv1.Query{{
    72  			QueryType: queryv1.QueryType_QUERY_TREE,
    73  			Tree: &queryv1.TreeQuery{
    74  				MaxNodes:           maxNodes,
    75  				StackTraceSelector: c.Msg.StackTraceSelector,
    76  				ProfileIdSelector:  c.Msg.ProfileIdSelector,
    77  			},
    78  		}},
    79  	})
    80  	if err != nil {
    81  		return nil, err
    82  	}
    83  	if report == nil {
    84  		return nil, nil
    85  	}
    86  	return report.Tree.Tree, nil
    87  }