github.com/grafana/pyroscope@v1.18.0/pkg/model/pprofsplit/pprof_split_by.go (about)

     1  package pprofsplit
     2  
     3  import (
     4  	"strings"
     5  
     6  	"github.com/prometheus/prometheus/model/relabel"
     7  
     8  	profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
     9  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    10  )
    11  
    12  // VisitSampleSeriesBy visits samples in a profile, splitting them by the
    13  // specified labels. Labels shared by all samples in a group (intersection)
    14  // are passed to the visitor, while the rest of the labels are added to each
    15  // sample.
    16  func VisitSampleSeriesBy(
    17  	profile *profilev1.Profile,
    18  	labels phlaremodel.Labels,
    19  	rules []*relabel.Config,
    20  	visitor SampleSeriesVisitor,
    21  	names ...string,
    22  ) error {
    23  	m := &sampleSeriesMerger{
    24  		visitor: visitor,
    25  		profile: profile,
    26  		names:   names,
    27  		groups:  make(map[string]*groupBy),
    28  	}
    29  	if err := VisitSampleSeries(profile, labels, rules, m); err != nil {
    30  		return err
    31  	}
    32  	if len(m.groups) == 0 {
    33  		return nil
    34  	}
    35  	for _, k := range m.order {
    36  		group := m.groups[k]
    37  		// For simplicity, we allocate a new slice of samples
    38  		// and delegate ownership to the visitor.
    39  		var size int
    40  		for _, s := range group.samples {
    41  			size += len(s.samples)
    42  		}
    43  		samples := make([]*profilev1.Sample, 0, size)
    44  		// All sample groups share group.labels:
    45  		// we use them as series labels.
    46  		for _, s := range group.samples {
    47  			s.labels = s.labels.Subtract(group.labels)
    48  			m.addLabelsToSamples(s)
    49  			samples = append(samples, s.samples...)
    50  		}
    51  		m.visitor.VisitSampleSeries(group.labels, samples)
    52  	}
    53  	return nil
    54  }
    55  
    56  type sampleSeriesMerger struct {
    57  	visitor SampleSeriesVisitor
    58  	profile *profilev1.Profile
    59  	names   []string
    60  	groups  map[string]*groupBy
    61  	order   []string
    62  	strings map[string]int64
    63  }
    64  
    65  type groupBy struct {
    66  	labels  phlaremodel.Labels
    67  	samples []sampleGroup
    68  	init    bool
    69  }
    70  
    71  type sampleGroup struct {
    72  	labels  phlaremodel.Labels
    73  	samples []*profilev1.Sample
    74  }
    75  
    76  func (m *sampleSeriesMerger) ValidateLabels(labels phlaremodel.Labels) (phlaremodel.Labels, error) {
    77  	return m.visitor.ValidateLabels(labels)
    78  }
    79  
    80  func (m *sampleSeriesMerger) VisitProfile(labels phlaremodel.Labels) {
    81  	m.visitor.VisitProfile(labels)
    82  }
    83  
    84  func (m *sampleSeriesMerger) VisitSampleSeries(labels phlaremodel.Labels, samples []*profilev1.Sample) {
    85  	k := groupKey(labels, m.names...)
    86  	group, ok := m.groups[k]
    87  	if !ok {
    88  		group = &groupBy{labels: labels.Clone()}
    89  		m.order = append(m.order, k)
    90  		m.groups[k] = group
    91  	} else {
    92  		group.labels = group.labels.Intersect(labels)
    93  	}
    94  	group.samples = append(group.samples, sampleGroup{
    95  		labels:  labels,
    96  		samples: samples,
    97  	})
    98  }
    99  
   100  func (m *sampleSeriesMerger) Discarded(profiles, bytes int) {
   101  	m.visitor.Discarded(profiles, bytes)
   102  }
   103  
   104  func groupKey(labels phlaremodel.Labels, by ...string) string {
   105  	var b strings.Builder
   106  	for _, name := range by {
   107  		for i := range labels {
   108  			if labels[i] != nil && labels[i].Name == name {
   109  				if b.Len() > 0 {
   110  					b.WriteByte(',')
   111  				}
   112  				b.WriteString(labels[i].Value)
   113  				break
   114  			}
   115  		}
   116  	}
   117  	return b.String()
   118  }
   119  
   120  func (m *sampleSeriesMerger) addLabelsToSamples(s sampleGroup) {
   121  	sampleLabels := make([]*profilev1.Label, len(s.labels))
   122  	for i, label := range s.labels {
   123  		sampleLabels[i] = &profilev1.Label{
   124  			Key: m.string(label.Name),
   125  			Str: m.string(label.Value),
   126  		}
   127  	}
   128  	// We can't reuse labels here.
   129  	for _, sample := range s.samples {
   130  		for i := range sampleLabels {
   131  			sample.Label = append(sample.Label, sampleLabels[i].CloneVT())
   132  		}
   133  	}
   134  }
   135  
   136  func (m *sampleSeriesMerger) string(s string) int64 {
   137  	if m.strings == nil {
   138  		m.strings = make(map[string]int64, len(m.profile.StringTable))
   139  		for i, str := range m.profile.StringTable {
   140  			m.strings[str] = int64(i)
   141  		}
   142  	}
   143  	i, ok := m.strings[s]
   144  	if !ok {
   145  		i = int64(len(m.profile.StringTable))
   146  		m.strings[s] = i
   147  		m.profile.StringTable = append(m.profile.StringTable, s)
   148  	}
   149  	return i
   150  }