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

     1  package queryfrontend
     2  
     3  import (
     4  	"context"
     5  
     6  	"connectrpc.com/connect"
     7  	"github.com/grafana/dskit/tenant"
     8  
     9  	profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
    10  	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
    11  	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
    12  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    13  	"github.com/grafana/pyroscope/pkg/pprof"
    14  	"github.com/grafana/pyroscope/pkg/validation"
    15  )
    16  
    17  func (q *QueryFrontend) SelectMergeProfile(
    18  	ctx context.Context,
    19  	c *connect.Request[querierv1.SelectMergeProfileRequest],
    20  ) (*connect.Response[profilev1.Profile], error) {
    21  	tenantIDs, err := tenant.TenantIDs(ctx)
    22  	if err != nil {
    23  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    24  	}
    25  	empty, err := validation.SanitizeTimeRange(q.limits, tenantIDs, &c.Msg.Start, &c.Msg.End)
    26  	if err != nil {
    27  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    28  	}
    29  	if empty {
    30  		return connect.NewResponse(&profilev1.Profile{}), nil
    31  	}
    32  
    33  	_, err = phlaremodel.ParseProfileTypeSelector(c.Msg.ProfileTypeID)
    34  	if err != nil {
    35  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    36  	}
    37  
    38  	// NOTE: Max nodes limit is off by default. SelectMergeProfile is primarily
    39  	// used for pprof export where truncation would result in incomplete profiles.
    40  	// This can be overridden per-tenant to enable enforcement if needed.
    41  	maxNodesEnabled := false
    42  	for _, tenantID := range tenantIDs {
    43  		if q.limits.MaxFlameGraphNodesOnSelectMergeProfile(tenantID) {
    44  			maxNodesEnabled = true
    45  		}
    46  	}
    47  
    48  	maxNodes := c.Msg.GetMaxNodes()
    49  	if maxNodesEnabled {
    50  		maxNodes, err = validation.ValidateMaxNodes(q.limits, tenantIDs, c.Msg.GetMaxNodes())
    51  		if err != nil {
    52  			return nil, connect.NewError(connect.CodeInvalidArgument, err)
    53  		}
    54  	}
    55  
    56  	labelSelector, err := buildLabelSelectorWithProfileType(c.Msg.LabelSelector, c.Msg.ProfileTypeID)
    57  	if err != nil {
    58  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    59  	}
    60  	report, err := q.querySingle(ctx, &queryv1.QueryRequest{
    61  		StartTime:     c.Msg.Start,
    62  		EndTime:       c.Msg.End,
    63  		LabelSelector: labelSelector,
    64  		Query: []*queryv1.Query{{
    65  			QueryType: queryv1.QueryType_QUERY_PPROF,
    66  			Pprof: &queryv1.PprofQuery{
    67  				MaxNodes:           maxNodes,
    68  				StackTraceSelector: c.Msg.StackTraceSelector,
    69  				ProfileIdSelector:  c.Msg.ProfileIdSelector,
    70  			},
    71  		}},
    72  	})
    73  	if err != nil {
    74  		return nil, err
    75  	}
    76  	if report == nil {
    77  		return connect.NewResponse(&profilev1.Profile{}), nil
    78  	}
    79  	var p profilev1.Profile
    80  	if err = pprof.Unmarshal(report.Pprof.Pprof, &p); err != nil {
    81  		return nil, err
    82  	}
    83  
    84  	return connect.NewResponse(&p), nil
    85  }