github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/convert/pprof/streaming/labels_cache.go (about)

     1  package streaming
     2  
     3  import (
     4  	"github.com/cespare/xxhash/v2"
     5  	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
     6  	"github.com/pyroscope-io/pyroscope/pkg/util/arenahelper"
     7  	"golang.org/x/exp/slices"
     8  	"reflect"
     9  	"unsafe"
    10  )
    11  
    12  type Labels []uint64
    13  
    14  func (l Labels) Len() int           { return len(l) }
    15  func (l Labels) Less(i, j int) bool { return l[i] < l[j] }
    16  func (l Labels) Swap(i, j int)      { l[i], l[j] = l[j], l[i] }
    17  
    18  var zeroHash = xxhash.Sum64(nil)
    19  
    20  func (l Labels) Hash() uint64 {
    21  	if len(l) == 0 {
    22  		return zeroHash
    23  	}
    24  	slices.Sort(l) //sort.Sort(l) // slice to interface conversion does an allocation for a slice header copy
    25  	var raw []byte
    26  	sh := (*reflect.SliceHeader)(unsafe.Pointer(&raw))
    27  	sh.Data = uintptr(unsafe.Pointer(&l[0]))
    28  	sh.Len = len(l) * 8
    29  	sh.Cap = len(l) * 8
    30  	return xxhash.Sum64(raw)
    31  }
    32  
    33  type LabelsCache struct {
    34  	// sample type -> labels hash -> index to labelRefs and trees
    35  	indices []map[uint64]int
    36  	// A label reference points to the subset in labels:
    37  	// Hight 32 bits is the start offset, lower 32 bits is the subset size.
    38  	labelRefs []uint64
    39  	labels    []uint64 // Packed label Key and Value indices
    40  	trees     []*tree.Tree
    41  
    42  	arena arenahelper.ArenaWrapper
    43  }
    44  
    45  func (c *LabelsCache) Reset() {
    46  	if c.indices == nil {
    47  		c.indices = arenahelper.MakeSlice[map[uint64]int](c.arena, 4, 4)
    48  	} else {
    49  		for i := range c.indices {
    50  			c.indices[i] = nil
    51  		}
    52  		for i := range c.trees {
    53  			c.trees[i] = nil
    54  		}
    55  		c.labelRefs = c.labelRefs[:0]
    56  		c.labels = c.labels[:0]
    57  		c.trees = c.trees[:0]
    58  	}
    59  }
    60  
    61  func (c *LabelsCache) GetOrCreateTree(sampleTypeIndex int, l Labels) *tree.Tree {
    62  	if sampleTypeIndex >= len(c.indices) {
    63  		newSampleTypes := make([]map[uint64]int, sampleTypeIndex+1, sampleTypeIndex+1)
    64  		copy(newSampleTypes, c.indices)
    65  		c.indices = newSampleTypes
    66  	}
    67  	p := c.indices[sampleTypeIndex]
    68  	if p == nil {
    69  		e, t := c.newCacheEntryA(l)
    70  		c.indices[sampleTypeIndex] = map[uint64]int{l.Hash(): e}
    71  		return t
    72  	}
    73  	h := l.Hash()
    74  	e, found := p[h]
    75  	if found {
    76  		return c.trees[e]
    77  	}
    78  	e, t := c.newCacheEntryA(l)
    79  	p[h] = e
    80  	return t
    81  }
    82  
    83  func (c *LabelsCache) Get(sampleTypeIndex int, h uint64) (*tree.Tree, bool) {
    84  	if sampleTypeIndex >= len(c.indices) {
    85  		return nil, false
    86  	}
    87  	p := c.indices[sampleTypeIndex]
    88  	if p == nil {
    89  		return nil, false
    90  	}
    91  	x, ok := p[h]
    92  	return c.trees[x], ok
    93  }
    94  
    95  func (c *LabelsCache) Remove(sampleTypeIndex int, h uint64) {
    96  	if sampleTypeIndex >= len(c.indices) {
    97  		return
    98  	}
    99  	p := c.indices[sampleTypeIndex]
   100  	if p == nil {
   101  		return
   102  	}
   103  	delete(p, h)
   104  	if len(p) == 0 {
   105  		c.indices[sampleTypeIndex] = nil
   106  	}
   107  }
   108  
   109  func (c *LabelsCache) newCacheEntry(l Labels) (int, *tree.Tree) {
   110  	from := len(c.labels)
   111  	for _, u := range l {
   112  		c.labels = append(c.labels, u)
   113  	}
   114  	to := len(c.labels)
   115  	res := len(c.labelRefs)
   116  	c.labelRefs = append(c.labelRefs, uint64(from<<32|to))
   117  	t := tree.New()
   118  	c.trees = append(c.trees, t)
   119  	return res, t
   120  }
   121  
   122  func (c *LabelsCache) iterate(f func(sampleTypeIndex int, l Labels, lh uint64, t *tree.Tree) error) error {
   123  	for sampleTypeIndex, p := range c.indices {
   124  		if p == nil {
   125  			continue
   126  		}
   127  		for h, x := range p {
   128  			labelRef := c.labelRefs[x]
   129  			l := c.labels[labelRef>>32 : labelRef&0xffffffff]
   130  			err := f(sampleTypeIndex, l, h, c.trees[x])
   131  			if err != nil {
   132  				return err
   133  			}
   134  		}
   135  	}
   136  	return nil
   137  }
   138  
   139  // CutLabel creates a copy of labels without label i.
   140  func CutLabel(a arenahelper.ArenaWrapper, labels Labels, i int) Labels {
   141  	c := arenahelper.MakeSlice[uint64](a, len(labels)-1, len(labels)-1)
   142  	copy(c[:i], labels[:i])
   143  	copy(c[i:], labels[i+1:])
   144  	return c
   145  }