github.com/grafana/pyroscope@v1.18.0/pkg/frontend/frontend_select_merge_span_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  	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) SelectMergeSpanProfile(
    21  	ctx context.Context,
    22  	c *connect.Request[querierv1.SelectMergeSpanProfileRequest],
    23  ) (*connect.Response[querierv1.SelectMergeSpanProfileResponse], error) {
    24  	ctx = connectgrpc.WithProcedure(ctx, querierv1connect.QuerierServiceSelectMergeSpanProfileProcedure)
    25  	tenantIDs, err := tenant.TenantIDs(ctx)
    26  	if err != nil {
    27  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    28  	}
    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(&querierv1.SelectMergeSpanProfileResponse{Flamegraph: &querierv1.FlameGraph{}}), nil
    36  	}
    37  	maxNodes, err := validation.ValidateMaxNodes(f.limits, tenantIDs, c.Msg.GetMaxNodes())
    38  	if err != nil {
    39  		return nil, connect.NewError(connect.CodeInvalidArgument, err)
    40  	}
    41  
    42  	g, ctx := errgroup.WithContext(ctx)
    43  	if maxConcurrent := validationutil.SmallestPositiveNonZeroIntPerTenant(tenantIDs, f.limits.MaxQueryParallelism); maxConcurrent > 0 {
    44  		g.SetLimit(maxConcurrent)
    45  	}
    46  
    47  	m := phlaremodel.NewFlameGraphMerger()
    48  	interval := validationutil.MaxDurationOrZeroPerTenant(tenantIDs, f.limits.QuerySplitDuration)
    49  	intervals := NewTimeIntervalIterator(time.UnixMilli(int64(validated.Start)), time.UnixMilli(int64(validated.End)), interval)
    50  
    51  	for intervals.Next() {
    52  		r := intervals.At()
    53  		g.Go(func() error {
    54  			req := connectgrpc.CloneRequest(c, &querierv1.SelectMergeSpanProfileRequest{
    55  				ProfileTypeID: c.Msg.ProfileTypeID,
    56  				LabelSelector: c.Msg.LabelSelector,
    57  				Start:         r.Start.UnixMilli(),
    58  				End:           r.End.UnixMilli(),
    59  				MaxNodes:      &maxNodes,
    60  				SpanSelector:  c.Msg.SpanSelector,
    61  				Format:        querierv1.ProfileFormat_PROFILE_FORMAT_TREE,
    62  			})
    63  			resp, err := connectgrpc.RoundTripUnary[
    64  				querierv1.SelectMergeSpanProfileRequest,
    65  				querierv1.SelectMergeSpanProfileResponse](ctx, f, req)
    66  			if err != nil {
    67  				return err
    68  			}
    69  			if len(resp.Msg.Tree) > 0 {
    70  				err = m.MergeTreeBytes(resp.Msg.Tree)
    71  			} else if resp.Msg.Flamegraph != nil {
    72  				// For backward compatibility.
    73  				m.MergeFlameGraph(resp.Msg.Flamegraph)
    74  			}
    75  			return err
    76  		})
    77  	}
    78  
    79  	if err = g.Wait(); err != nil {
    80  		return nil, err
    81  	}
    82  
    83  	t := m.Tree()
    84  	var resp querierv1.SelectMergeSpanProfileResponse
    85  	switch c.Msg.Format {
    86  	default:
    87  		resp.Flamegraph = phlaremodel.NewFlameGraph(t, c.Msg.GetMaxNodes())
    88  	case querierv1.ProfileFormat_PROFILE_FORMAT_TREE:
    89  		resp.Tree = t.Bytes(c.Msg.GetMaxNodes())
    90  	}
    91  	return connect.NewResponse(&resp), nil
    92  }