github.com/grafana/pyroscope@v1.18.0/pkg/frontend/frontend_select_merge_stacktraces.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  	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
    13  	"github.com/grafana/pyroscope/api/gen/proto/go/querier/v1/querierv1connect"
    14  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    15  	"github.com/grafana/pyroscope/pkg/util/connectgrpc"
    16  	validationutil "github.com/grafana/pyroscope/pkg/util/validation"
    17  	"github.com/grafana/pyroscope/pkg/validation"
    18  )
    19  
    20  func (f *Frontend) SelectMergeStacktraces(
    21  	ctx context.Context,
    22  	c *connect.Request[querierv1.SelectMergeStacktracesRequest],
    23  ) (*connect.Response[querierv1.SelectMergeStacktracesResponse], error) {
    24  	t, err := f.selectMergeStacktracesTree(ctx, c)
    25  	if err != nil {
    26  		return nil, err
    27  	}
    28  	var resp querierv1.SelectMergeStacktracesResponse
    29  	switch c.Msg.Format {
    30  	default:
    31  		resp.Flamegraph = phlaremodel.NewFlameGraph(t, c.Msg.GetMaxNodes())
    32  	case querierv1.ProfileFormat_PROFILE_FORMAT_TREE:
    33  		resp.Tree = t.Bytes(c.Msg.GetMaxNodes())
    34  	}
    35  	return connect.NewResponse(&resp), nil
    36  }
    37  
    38  func (f *Frontend) selectMergeStacktracesTree(
    39  	ctx context.Context,
    40  	c *connect.Request[querierv1.SelectMergeStacktracesRequest],
    41  ) (*phlaremodel.Tree, error) {
    42  	ctx = connectgrpc.WithProcedure(ctx, querierv1connect.QuerierServiceSelectMergeStacktracesProcedure)
    43  	tenantIDs, err := tenant.TenantIDs(ctx)
    44  	if err != nil {
    45  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    46  	}
    47  
    48  	validated, err := validation.ValidateRangeRequest(f.limits, tenantIDs, model.Interval{Start: model.Time(c.Msg.Start), End: model.Time(c.Msg.End)}, model.Now())
    49  	if err != nil {
    50  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    51  	}
    52  	if validated.IsEmpty {
    53  		return new(phlaremodel.Tree), nil
    54  	}
    55  	maxNodes, err := validation.ValidateMaxNodes(f.limits, tenantIDs, c.Msg.GetMaxNodes())
    56  	if err != nil {
    57  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    58  	}
    59  
    60  	g, ctx := errgroup.WithContext(ctx)
    61  	if maxConcurrent := validationutil.SmallestPositiveNonZeroIntPerTenant(tenantIDs, f.limits.MaxQueryParallelism); maxConcurrent > 0 {
    62  		g.SetLimit(maxConcurrent)
    63  	}
    64  
    65  	m := phlaremodel.NewFlameGraphMerger()
    66  	interval := validationutil.MaxDurationOrZeroPerTenant(tenantIDs, f.limits.QuerySplitDuration)
    67  	intervals := NewTimeIntervalIterator(time.UnixMilli(int64(validated.Start)), time.UnixMilli(int64(validated.End)), interval)
    68  
    69  	for intervals.Next() {
    70  		r := intervals.At()
    71  		g.Go(func() error {
    72  			req := connectgrpc.CloneRequest(c, &querierv1.SelectMergeStacktracesRequest{
    73  				ProfileTypeID: c.Msg.ProfileTypeID,
    74  				LabelSelector: c.Msg.LabelSelector,
    75  				Start:         r.Start.UnixMilli(),
    76  				End:           r.End.UnixMilli(),
    77  				MaxNodes:      &maxNodes,
    78  				Format:        querierv1.ProfileFormat_PROFILE_FORMAT_TREE,
    79  			})
    80  			resp, err := connectgrpc.RoundTripUnary[
    81  				querierv1.SelectMergeStacktracesRequest,
    82  				querierv1.SelectMergeStacktracesResponse](ctx, f, req)
    83  			if err != nil {
    84  				return err
    85  			}
    86  			if len(resp.Msg.Tree) > 0 {
    87  				err = m.MergeTreeBytes(resp.Msg.Tree)
    88  			} else if resp.Msg.Flamegraph != nil {
    89  				// For backward compatibility.
    90  				m.MergeFlameGraph(resp.Msg.Flamegraph)
    91  			}
    92  			return err
    93  		})
    94  	}
    95  
    96  	if err = g.Wait(); err != nil {
    97  		return nil, err
    98  	}
    99  
   100  	return m.Tree(), nil
   101  }