github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/filter_profiles_bidi.go (about)

     1  package phlaredb
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  
     7  	"connectrpc.com/connect"
     8  	"github.com/opentracing/opentracing-go"
     9  	otlog "github.com/opentracing/opentracing-go/log"
    10  	"github.com/pkg/errors"
    11  	"github.com/prometheus/common/model"
    12  
    13  	ingestv1 "github.com/grafana/pyroscope/api/gen/proto/go/ingester/v1"
    14  	"github.com/grafana/pyroscope/pkg/iter"
    15  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    16  )
    17  
    18  type BidiServerMerge[Res any, Req any] interface {
    19  	Send(Res) error
    20  	Receive() (Req, error)
    21  }
    22  
    23  type ProfileWithIndex struct {
    24  	Profile
    25  	Index int
    26  }
    27  
    28  type indexedProfileIterator struct {
    29  	iter.Iterator[Profile]
    30  	querierIndex int
    31  }
    32  
    33  func (pqi *indexedProfileIterator) At() ProfileWithIndex {
    34  	return ProfileWithIndex{
    35  		Profile: pqi.Iterator.At(),
    36  		Index:   pqi.querierIndex,
    37  	}
    38  }
    39  
    40  type filterRequest interface {
    41  	*ingestv1.MergeProfilesStacktracesRequest |
    42  		*ingestv1.MergeProfilesLabelsRequest |
    43  		*ingestv1.MergeProfilesPprofRequest |
    44  		*ingestv1.MergeSpanProfileRequest
    45  }
    46  
    47  type filterResponse interface {
    48  	*ingestv1.MergeProfilesStacktracesResponse |
    49  		*ingestv1.MergeProfilesLabelsResponse |
    50  		*ingestv1.MergeProfilesPprofResponse |
    51  		*ingestv1.MergeSpanProfileResponse
    52  }
    53  
    54  func rewriteEOFError(err error) error {
    55  	if errors.Is(err, io.EOF) {
    56  		return connect.NewError(connect.CodeCanceled, errors.New("client closed stream"))
    57  	}
    58  	return err
    59  }
    60  
    61  // filterProfiles merges and dedupe profiles from different iterators and allow filtering via a bidi stream.
    62  func filterProfiles[B BidiServerMerge[Res, Req], Res filterResponse, Req filterRequest](
    63  	ctx context.Context, profiles []iter.Iterator[Profile], batchProfileSize int, stream B,
    64  ) ([][]Profile, error) {
    65  	sp, _ := opentracing.StartSpanFromContext(ctx, "filterProfiles")
    66  	defer sp.Finish()
    67  	selection := make([][]Profile, len(profiles))
    68  	selectProfileResult := &ingestv1.ProfileSets{
    69  		Profiles:     make([]*ingestv1.SeriesProfile, 0, batchProfileSize),
    70  		Fingerprints: make([]uint64, 0, batchProfileSize),
    71  	}
    72  	its := make([]iter.Iterator[ProfileWithIndex], len(profiles))
    73  	for i, iter := range profiles {
    74  		iter := iter
    75  		its[i] = &indexedProfileIterator{
    76  			Iterator:     iter,
    77  			querierIndex: i,
    78  		}
    79  	}
    80  	if err := iter.ReadBatch(ctx, phlaremodel.NewMergeIterator(ProfileWithIndex{
    81  		Profile: maxBlockProfile,
    82  		Index:   0,
    83  	}, true, its...), batchProfileSize, func(ctx context.Context, batch []ProfileWithIndex) error {
    84  		sp.LogFields(
    85  			otlog.Int("batch_len", len(batch)),
    86  			otlog.Int("batch_requested_size", batchProfileSize),
    87  		)
    88  
    89  		seriesByFP := map[model.Fingerprint]int{}
    90  		selectProfileResult.Profiles = selectProfileResult.Profiles[:0]
    91  		selectProfileResult.Fingerprints = selectProfileResult.Fingerprints[:0]
    92  
    93  		for _, profile := range batch {
    94  			var ok bool
    95  			var idx int
    96  			fp := profile.Fingerprint()
    97  			idx, ok = seriesByFP[fp]
    98  			if !ok {
    99  				idx = len(selectProfileResult.Fingerprints)
   100  				seriesByFP[fp] = idx
   101  				selectProfileResult.Fingerprints = append(selectProfileResult.Fingerprints, uint64(fp))
   102  			}
   103  			selectProfileResult.Profiles = append(selectProfileResult.Profiles, &ingestv1.SeriesProfile{
   104  				LabelIndex: int32(idx),
   105  				Timestamp:  int64(profile.Timestamp()),
   106  			})
   107  		}
   108  
   109  		sp.LogFields(otlog.String("msg", "sending batch to client"))
   110  		var err error
   111  		switch s := BidiServerMerge[Res, Req](stream).(type) {
   112  		case BidiServerMerge[*ingestv1.MergeProfilesStacktracesResponse, *ingestv1.MergeProfilesStacktracesRequest]:
   113  			err = s.Send(&ingestv1.MergeProfilesStacktracesResponse{
   114  				SelectedProfiles: selectProfileResult,
   115  			})
   116  		case BidiServerMerge[*ingestv1.MergeProfilesLabelsResponse, *ingestv1.MergeProfilesLabelsRequest]:
   117  			err = s.Send(&ingestv1.MergeProfilesLabelsResponse{
   118  				SelectedProfiles: selectProfileResult,
   119  			})
   120  		case BidiServerMerge[*ingestv1.MergeProfilesPprofResponse, *ingestv1.MergeProfilesPprofRequest]:
   121  			err = s.Send(&ingestv1.MergeProfilesPprofResponse{
   122  				SelectedProfiles: selectProfileResult,
   123  			})
   124  		case BidiServerMerge[*ingestv1.MergeSpanProfileResponse, *ingestv1.MergeSpanProfileRequest]:
   125  			err = s.Send(&ingestv1.MergeSpanProfileResponse{
   126  				SelectedProfiles: selectProfileResult,
   127  			})
   128  		}
   129  		// read a batch of profiles and sends it.
   130  
   131  		if err != nil {
   132  			return rewriteEOFError(err)
   133  		}
   134  		sp.LogFields(otlog.String("msg", "batch sent to client"))
   135  
   136  		sp.LogFields(otlog.String("msg", "reading selection from client"))
   137  
   138  		// handle response for the batch.
   139  		var selected []bool
   140  		switch s := BidiServerMerge[Res, Req](stream).(type) {
   141  		case BidiServerMerge[*ingestv1.MergeProfilesStacktracesResponse, *ingestv1.MergeProfilesStacktracesRequest]:
   142  			selectionResponse, err := s.Receive()
   143  			if err != nil {
   144  				return rewriteEOFError(err)
   145  			}
   146  			selected = selectionResponse.Profiles
   147  		case BidiServerMerge[*ingestv1.MergeProfilesLabelsResponse, *ingestv1.MergeProfilesLabelsRequest]:
   148  			selectionResponse, err := s.Receive()
   149  			if err != nil {
   150  				return rewriteEOFError(err)
   151  			}
   152  			selected = selectionResponse.Profiles
   153  		case BidiServerMerge[*ingestv1.MergeProfilesPprofResponse, *ingestv1.MergeProfilesPprofRequest]:
   154  			selectionResponse, err := s.Receive()
   155  			if err != nil {
   156  				return rewriteEOFError(err)
   157  			}
   158  			selected = selectionResponse.Profiles
   159  		case BidiServerMerge[*ingestv1.MergeSpanProfileResponse, *ingestv1.MergeSpanProfileRequest]:
   160  			selectionResponse, err := s.Receive()
   161  			if err != nil {
   162  				return rewriteEOFError(err)
   163  			}
   164  			selected = selectionResponse.Profiles
   165  		}
   166  		sp.LogFields(otlog.String("msg", "selection received"))
   167  		for i, k := range selected {
   168  			if k {
   169  				selection[batch[i].Index] = append(selection[batch[i].Index], batch[i].Profile)
   170  			}
   171  		}
   172  		return nil
   173  	}); err != nil {
   174  		return nil, err
   175  	}
   176  	return selection, nil
   177  }