github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/storage/dict/dict.go (about)

     1  package dict
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/pyroscope-io/pyroscope/pkg/util/varint"
     9  	"github.com/valyala/bytebufferpool"
    10  )
    11  
    12  type (
    13  	Key   []byte
    14  	Value []byte
    15  )
    16  
    17  func New() *Dict {
    18  	return &Dict{
    19  		root: newTrieNode([]byte{}),
    20  	}
    21  }
    22  
    23  type Dict struct {
    24  	m    sync.RWMutex
    25  	root *trieNode
    26  }
    27  
    28  func (t *Dict) GetValue(key Key, value io.Writer) bool {
    29  	t.m.RLock()
    30  	defer t.m.RUnlock()
    31  	return t.readValue(key, value)
    32  }
    33  
    34  func (t *Dict) Get(key Key) (Value, bool) {
    35  	t.m.RLock()
    36  	defer t.m.RUnlock()
    37  	var labelBuf bytes.Buffer
    38  	if t.readValue(key, &labelBuf) {
    39  		return labelBuf.Bytes(), true
    40  	}
    41  	return nil, false
    42  }
    43  
    44  func (t *Dict) readValue(key Key, w io.Writer) bool {
    45  	r := bytes.NewReader(key)
    46  	tn := t.root
    47  	for {
    48  		v, err := varint.Read(r)
    49  		if err != nil {
    50  			return true
    51  		}
    52  		if int(v) >= len(tn.children) {
    53  			return false
    54  		}
    55  
    56  		label := tn.children[v].label
    57  		_, _ = w.Write(label)
    58  		tn = tn.children[v]
    59  
    60  		expectedLen, _ := varint.Read(r)
    61  		for len(label) < int(expectedLen) {
    62  			if len(tn.children) == 0 {
    63  				return false
    64  			}
    65  			label2 := tn.children[0].label
    66  			_, _ = w.Write(label2)
    67  			expectedLen -= uint64(len(label2))
    68  			tn = tn.children[0]
    69  		}
    70  	}
    71  }
    72  
    73  var writerPool = sync.Pool{New: func() any { return varint.NewWriter() }}
    74  
    75  func (t *Dict) PutValue(val Value, dst io.Writer) {
    76  	t.m.Lock()
    77  	defer t.m.Unlock()
    78  	vw := writerPool.Get().(varint.Writer)
    79  	defer writerPool.Put(vw)
    80  	t.root.findNodeAt(val, vw, dst)
    81  }
    82  
    83  var bufferPool bytebufferpool.Pool
    84  
    85  func (t *Dict) Put(val Value) Key {
    86  	b := bufferPool.Get()
    87  	defer bufferPool.Put(b)
    88  	t.PutValue(val, b)
    89  	k := make([]byte, b.Len())
    90  	copy(k, b.B)
    91  	return k
    92  }