github.com/grafana/pyroscope@v1.18.0/pkg/og/storage/tree/profile_extra.go (about)

     1  package tree
     2  
     3  // These functions are kept separately as profile.pb.go is a generated file
     4  
     5  import (
     6  	"sort"
     7  
     8  	profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
     9  	"github.com/grafana/pyroscope/pkg/og/agent/spy"
    10  	"github.com/valyala/bytebufferpool"
    11  )
    12  
    13  type cacheKey []int64
    14  
    15  type cacheEntry struct {
    16  	key cacheKey
    17  	val *spy.Labels
    18  }
    19  type cache struct {
    20  	data []*cacheEntry
    21  }
    22  
    23  func newCache() *cache {
    24  	return &cache{
    25  		data: []*cacheEntry{},
    26  	}
    27  }
    28  
    29  func getCacheKey(l []*profilev1.Label) cacheKey {
    30  	r := []int64{}
    31  	for _, x := range l {
    32  		if x.Str != 0 {
    33  			r = append(r, x.Key, x.Str)
    34  		}
    35  	}
    36  	return r
    37  }
    38  
    39  func eq(a, b []int64) bool {
    40  	if len(a) != len(b) {
    41  		return false
    42  	}
    43  	for i, v := range a {
    44  		if v != b[i] {
    45  			return false
    46  		}
    47  	}
    48  	return true
    49  }
    50  
    51  func (c *cache) pprofLabelsToSpyLabels(x *profilev1.Profile, pprofLabels []*profilev1.Label) *spy.Labels {
    52  	k := getCacheKey(pprofLabels)
    53  	for _, e := range c.data {
    54  		if eq(e.key, k) {
    55  			return e.val
    56  		}
    57  	}
    58  
    59  	l := spy.NewLabels()
    60  	for _, pl := range pprofLabels {
    61  		if pl.Str != 0 {
    62  			l.Set(x.StringTable[pl.Key], x.StringTable[pl.Str])
    63  		}
    64  	}
    65  	newVal := &cacheEntry{
    66  		key: k,
    67  		val: l,
    68  	}
    69  	c.data = append(c.data, newVal)
    70  	return l
    71  }
    72  
    73  func Get(x *profilev1.Profile, sampleType string, cb func(labels *spy.Labels, name []byte, val int) error) error {
    74  	valueIndex := 0
    75  	if sampleType != "" {
    76  		for i, v := range x.SampleType {
    77  			if x.StringTable[v.Type] == sampleType {
    78  				valueIndex = i
    79  				break
    80  			}
    81  		}
    82  	}
    83  
    84  	labelsCache := newCache()
    85  
    86  	b := bytebufferpool.Get()
    87  	defer bytebufferpool.Put(b)
    88  
    89  	for _, s := range x.Sample {
    90  		for i := len(s.LocationId) - 1; i >= 0; i-- {
    91  			name, ok := FindFunctionName(x, s.LocationId[i])
    92  			if !ok {
    93  				continue
    94  			}
    95  			if b.Len() > 0 {
    96  				_ = b.WriteByte(';')
    97  			}
    98  			_, _ = b.WriteString(name)
    99  		}
   100  
   101  		labels := labelsCache.pprofLabelsToSpyLabels(x, s.Label)
   102  		if err := cb(labels, b.Bytes(), int(s.Value[valueIndex])); err != nil {
   103  			return err
   104  		}
   105  
   106  		b.Reset()
   107  	}
   108  
   109  	return nil
   110  }
   111  
   112  func SampleTypes(x *profilev1.Profile) []string {
   113  	r := []string{}
   114  	for _, v := range x.SampleType {
   115  		r = append(r, x.StringTable[v.Type])
   116  	}
   117  	return r
   118  }
   119  
   120  func FindFunctionName(x *profilev1.Profile, locID uint64) (string, bool) {
   121  	if loc, ok := FindLocation(x, locID); ok {
   122  		if len(loc.Line) <= 0 {
   123  			return "", false
   124  		}
   125  
   126  		if fn, ok := FindFunction(x, loc.Line[0].FunctionId); ok {
   127  			return x.StringTable[fn.Name], true
   128  		}
   129  	}
   130  	return "", false
   131  }
   132  
   133  func FindLocation(x *profilev1.Profile, lid uint64) (*profilev1.Location, bool) {
   134  	idx := sort.Search(len(x.Location), func(i int) bool {
   135  		return x.Location[i].Id >= lid
   136  	})
   137  	if idx < len(x.Location) {
   138  		if l := x.Location[idx]; l.Id == lid {
   139  			return l, true
   140  		}
   141  	}
   142  	return nil, false
   143  }
   144  
   145  func FindFunction(x *profilev1.Profile, fid uint64) (*profilev1.Function, bool) {
   146  	idx := sort.Search(len(x.Function), func(i int) bool {
   147  		return x.Function[i].Id >= fid
   148  	})
   149  	if idx < len(x.Function) {
   150  		if f := x.Function[idx]; f.Id == fid {
   151  			return f, true
   152  		}
   153  	}
   154  	return nil, false
   155  }