perf: optimize allocations in dictionaries (#1610)

pkg/storage/dict/dict.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	"github.com/pyroscope-io/pyroscope/pkg/util/varint"
+	"github.com/valyala/bytebufferpool"
 )
 
 type (
@@ -69,10 +70,23 @@ func (t *Dict) readValue(key Key, w io.Writer) bool {
 	}
 }
 
-func (t *Dict) Put(val Value) Key {
+var writerPool = sync.Pool{New: func() any { return varint.NewWriter() }}
+
+func (t *Dict) PutValue(val Value, dst io.Writer) {
 	t.m.Lock()
 	defer t.m.Unlock()
-	var buf bytes.Buffer
-	t.root.findNodeAt(val, &buf)
-	return buf.Bytes()
+	vw := writerPool.Get().(varint.Writer)
+	defer writerPool.Put(vw)
+	t.root.findNodeAt(val, vw, dst)
+}
+
+var bufferPool bytebufferpool.Pool
+
+func (t *Dict) Put(val Value) Key {
+	b := bufferPool.Get()
+	defer bufferPool.Put(b)
+	t.PutValue(val, b)
+	k := make([]byte, b.Len())
+	copy(k, b.B)
+	return k
 }

pkg/storage/dict/trie.go

@@ -27,13 +27,12 @@ func (tn *trieNode) insert(t2 *trieNode) {
 }
 
 // TODO: too complicated, need to refactor / document this
-func (tn *trieNode) findNodeAt(key []byte, w io.Writer) {
+func (tn *trieNode) findNodeAt(key []byte, vw varint.Writer, w io.Writer) {
 	// log.Debug("findNodeAt")
 	key2 := make([]byte, len(key))
 	// TODO: remove
 	copy(key2, key)
 	key = key2
-	vw := varint.NewWriter()
 
 OuterLoop:
 	for {

pkg/storage/tree/serialize.go

@@ -25,18 +25,20 @@ func (t *Tree) SerializeTruncate(d *dict.Dict, maxNodes int, w io.Writer) error
 		return err
 	}
 
+	var b bytes.Buffer // Temporary buffer for dictionary keys.
 	minVal := t.minValue(maxNodes)
 	nodes := make([]*treeNode, 1, 128)
 	nodes[0] = t.root
 	for len(nodes) > 0 {
 		tn := nodes[0]
 		nodes = nodes[1:]
 
-		labelKey := d.Put([]byte(tn.Name))
-		if _, err = vw.Write(w, uint64(len(labelKey))); err != nil {
+		b.Reset()
+		d.PutValue([]byte(tn.Name), &b)
+		if _, err = vw.Write(w, uint64(b.Len())); err != nil {
 			return err
 		}
-		if _, err = w.Write(labelKey); err != nil {
+		if _, err = w.Write(b.Bytes()); err != nil {
 			return err
 		}
 		if _, err = vw.Write(w, tn.Self); err != nil {