github.com/grafana/pyroscope@v1.18.0/pkg/frontend/frontend_select_merge_profile.go (about)

     1  package frontend
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"connectrpc.com/connect"
     8  	"github.com/grafana/dskit/tenant"
     9  	"github.com/prometheus/common/model"
    10  	"golang.org/x/sync/errgroup"
    11  
    12  	profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
    13  	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
    14  	"github.com/grafana/pyroscope/api/gen/proto/go/querier/v1/querierv1connect"
    15  	"github.com/grafana/pyroscope/pkg/pprof"
    16  	"github.com/grafana/pyroscope/pkg/util/connectgrpc"
    17  	validationutil "github.com/grafana/pyroscope/pkg/util/validation"
    18  	"github.com/grafana/pyroscope/pkg/validation"
    19  )
    20  
    21  func (f *Frontend) SelectMergeProfile(
    22  	ctx context.Context,
    23  	c *connect.Request[querierv1.SelectMergeProfileRequest],
    24  ) (*connect.Response[profilev1.Profile], error) {
    25  	ctx = connectgrpc.WithProcedure(ctx, querierv1connect.QuerierServiceSelectMergeProfileProcedure)
    26  	tenantIDs, err := tenant.TenantIDs(ctx)
    27  	if err != nil {
    28  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    29  	}
    30  	validated, err := validation.ValidateRangeRequest(f.limits, tenantIDs, model.Interval{Start: model.Time(c.Msg.Start), End: model.Time(c.Msg.End)}, model.Now())
    31  	if err != nil {
    32  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    33  	}
    34  	if validated.IsEmpty {
    35  		return connect.NewResponse(&profilev1.Profile{}), nil
    36  	}
    37  	c.Msg.Start = int64(validated.Start)
    38  	c.Msg.End = int64(validated.End)
    39  
    40  	g, ctx := errgroup.WithContext(ctx)
    41  	if maxConcurrent := validationutil.SmallestPositiveNonZeroIntPerTenant(tenantIDs, f.limits.MaxQueryParallelism); maxConcurrent > 0 {
    42  		g.SetLimit(maxConcurrent)
    43  	}
    44  
    45  	interval := validationutil.MaxDurationOrZeroPerTenant(tenantIDs, f.limits.QuerySplitDuration)
    46  	intervals := NewTimeIntervalIterator(time.UnixMilli(int64(validated.Start)), time.UnixMilli(int64(validated.End)), interval)
    47  
    48  	// NOTE: Max nodes limit is off by default. SelectMergeProfile is primarily
    49  	// used for pprof export where truncation would result in incomplete profiles.
    50  	// This can be overridden per-tenant to enable enforcement if needed.
    51  	maxNodesEnabled := false
    52  	for _, tenantID := range tenantIDs {
    53  		if f.limits.MaxFlameGraphNodesOnSelectMergeProfile(tenantID) {
    54  			maxNodesEnabled = true
    55  		}
    56  	}
    57  
    58  	maxNodes := c.Msg.MaxNodes
    59  	if maxNodesEnabled {
    60  		maxNodesV, err := validation.ValidateMaxNodes(f.limits, tenantIDs, c.Msg.GetMaxNodes())
    61  		if err != nil {
    62  			return nil, connect.NewError(connect.CodeInvalidArgument, err)
    63  		}
    64  		maxNodes = &maxNodesV
    65  	}
    66  
    67  	var m pprof.ProfileMerge
    68  	for intervals.Next() {
    69  		r := intervals.At()
    70  		g.Go(func() error {
    71  			req := connectgrpc.CloneRequest(c, &querierv1.SelectMergeProfileRequest{
    72  				ProfileTypeID:      c.Msg.ProfileTypeID,
    73  				LabelSelector:      c.Msg.LabelSelector,
    74  				Start:              r.Start.UnixMilli(),
    75  				End:                r.End.UnixMilli(),
    76  				MaxNodes:           maxNodes,
    77  				StackTraceSelector: c.Msg.StackTraceSelector,
    78  			})
    79  			resp, err := connectgrpc.RoundTripUnary[
    80  				querierv1.SelectMergeProfileRequest,
    81  				profilev1.Profile](ctx, f, req)
    82  			if err != nil {
    83  				return err
    84  			}
    85  			return m.Merge(resp.Msg, f.limits.QuerySanitizeOnMerge(tenantIDs[0]))
    86  		})
    87  	}
    88  
    89  	if err = g.Wait(); err != nil {
    90  		return nil, err
    91  	}
    92  
    93  	return connect.NewResponse(m.Profile()), nil
    94  }