perf: optimize segment key parser allocations (#1625)

* perf: optimize segment key parser allocations

* perf: replace strings builder with bytes buffer

pkg/storage/segment/key.go

@@ -1,9 +1,11 @@
 package segment
 
 import (
+	"bytes"
 	"errors"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/pyroscope-io/pyroscope/pkg/flameql"
@@ -33,7 +35,9 @@ func NewKey(labels map[string]string) *Key { return &Key{labels: labels} }
 
 func ParseKey(name string) (*Key, error) {
 	k := &Key{labels: make(map[string]string)}
-	p := parser{parserState: nameParserState}
+	p := parserPool.Get().(*parser)
+	defer parserPool.Put(p)
+	p.reset()
 	var err error
 	for _, r := range name + "{" {
 		switch p.parserState {
@@ -53,54 +57,70 @@ func ParseKey(name string) (*Key, error) {
 
 type parser struct {
 	parserState ParserState
-	key         string
-	value       string
+	key         *bytes.Buffer
+	value       *bytes.Buffer
+}
+
+var parserPool = sync.Pool{
+	New: func() any {
+		return &parser{
+			parserState: nameParserState,
+			key:         new(bytes.Buffer),
+			value:       new(bytes.Buffer),
+		}
+	},
+}
+
+func (p *parser) reset() {
+	p.parserState = nameParserState
+	p.key.Reset()
+	p.value.Reset()
 }
 
 // ParseKey's nameParserState switch case
 func (p *parser) nameParserCase(r int32, k *Key) error {
 	switch r {
 	case '{':
 		p.parserState = tagKeyParserState
-		appName := strings.TrimSpace(p.value)
+		appName := strings.TrimSpace(p.value.String())
 		if err := flameql.ValidateAppName(appName); err != nil {
 			return err
 		}
 		k.labels["__name__"] = appName
 	default:
-		p.value += string(r)
+		p.value.WriteRune(r)
 	}
 	return nil
 }
 
 // ParseKey's tagKeyParserState switch case
-func (p *parser) tagKeyParserCase(r int32) {
+func (p *parser) tagKeyParserCase(r rune) {
 	switch r {
 	case '}':
 		p.parserState = doneParserState
 	case '=':
 		p.parserState = tagValueParserState
-		p.value = ""
+		p.value.Reset()
 	default:
-		p.key += string(r)
+		p.key.WriteRune(r)
 	}
 }
 
 // ParseKey's tagValueParserState switch case
-func (p *parser) tagValueParserCase(r int32, k *Key) error {
+func (p *parser) tagValueParserCase(r rune, k *Key) error {
 	switch r {
 	case ',', '}':
 		p.parserState = tagKeyParserState
-		key := strings.TrimSpace(p.key)
+		key := strings.TrimSpace(p.key.String())
 		if !flameql.IsTagKeyReserved(key) {
 			if err := flameql.ValidateTagKey(key); err != nil {
 				return err
 			}
 		}
-		k.labels[key] = strings.TrimSpace(p.value)
-		p.key = ""
+		k.labels[key] = strings.TrimSpace(p.value.String())
+		p.key.Reset()
 	default:
-		p.value += string(r)
+		p.value.WriteRune(r)
 	}
 	return nil
 }

pkg/storage/segment/key_bech_test.go

@@ -0,0 +1,45 @@
+package segment
+
+import (
+	"math/rand"
+	"testing"
+)
+
+func BenchmarkKey_Parse(b *testing.B) {
+	const (
+		labelsSize = 10
+		minLen     = 6
+		maxLen     = 16
+	)
+
+	// Duplicates are okay.
+	labels := make(map[string]string, labelsSize+1)
+	for i := 0; i < labelsSize; i++ {
+		labels[randString(randInt(minLen, maxLen))] = randString(randInt(minLen, maxLen))
+	}
+
+	labels["__name__"] = "benchmark.key.parse"
+	keyStr := NewKey(labels).Normalized()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		if _, err := ParseKey(keyStr); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// TODO(kolesnikovae): This is not near perfect way of generating strings.
+//  It makes sense to create a package for util functions like this.
+
+const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+
+func randString(n int) string {
+	b := make([]byte, n)
+	for i := range b {
+		b[i] = letterBytes[rand.Intn(len(letterBytes))]
+	}
+	return string(b)
+}