From f8d52f5d478024ad52107a6e1c3510b2a8fabdd6 Mon Sep 17 00:00:00 2001
From: Anton Kolesnikov <anton.e.kolesnikov@gmail.com>
Date: Tue, 18 Oct 2022 16:28:56 +0200
Subject: [PATCH] perf: optimize allocations in dictionaries (#1610)

---
 pkg/storage/dict/dict.go      | 22 ++++++++++++++++++----
 pkg/storage/dict/trie.go      |  3 +--
 pkg/storage/tree/serialize.go |  8 +++++---
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/pkg/storage/dict/dict.go b/pkg/storage/dict/dict.go
index 5c56321872..f8032c382e 100644
--- a/pkg/storage/dict/dict.go
+++ b/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
 }
diff --git a/pkg/storage/dict/trie.go b/pkg/storage/dict/trie.go
index ff24557340..0465f26555 100644
--- a/pkg/storage/dict/trie.go
+++ b/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 {
diff --git a/pkg/storage/tree/serialize.go b/pkg/storage/tree/serialize.go
index 88385de60f..0cf7edcc1b 100644
--- a/pkg/storage/tree/serialize.go
+++ b/pkg/storage/tree/serialize.go
@@ -25,6 +25,7 @@ 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
@@ -32,11 +33,12 @@ func (t *Tree) SerializeTruncate(d *dict.Dict, maxNodes int, w io.Writer) error
 		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 {