Use byte buffer pool for serialization

grafana · Jul 26, 2021 · 725cfdd · 725cfdd
1 parent 46b5aa1
commit 725cfdd
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 6 deletions.
diff --git a/go.mod b/go.mod
@@ -46,6 +46,7 @@ require (
 	github.com/sirupsen/logrus v1.7.0
 	github.com/tklauser/go-sysconf v0.3.6 // indirect
 	github.com/twmb/murmur3 v1.1.5
+	github.com/valyala/bytebufferpool v1.0.0
 	github.com/wacul/ptr v1.0.0 // indirect
 	golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421
 	golang.org/x/sync v0.0.0-20201207232520-09787c993a3a

diff --git a/go.sum b/go.sum
@@ -439,6 +439,8 @@ github.com/twmb/murmur3 v1.1.5/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq
 github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8/go.mod h1:VFNgLljTbGfSG7qAOspJ7OScBnGdDN/yBr0sguwnwf0=
 github.com/urfave/cli v1.20.0/go.mod h1:70zkFmudgCuE/ngEzBv17Jvp/497gISqfk5gWijbERA=
 github.com/urfave/cli v1.22.1/go.mod h1:Gos4lmkARVdJ6EkW0WaNv/tZAAMe9V7XWyB60NtXRu0=
+github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
+github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 github.com/wacul/ptr v0.0.0-20170209030335-91632201dfc8/go.mod h1:BD0gjsZrCwtoR+yWDB9v2hQ8STlq9tT84qKfa+3txOc=
 github.com/wacul/ptr v1.0.0 h1:FIKu08Wx0YUIf9MNsfF62OCmBSmz5A1Tk65zWhOIL/I=
 github.com/wacul/ptr v1.0.0/go.mod h1:BD0gjsZrCwtoR+yWDB9v2hQ8STlq9tT84qKfa+3txOc=

diff --git a/pkg/storage/cache/cache.go b/pkg/storage/cache/cache.go
@@ -3,12 +3,13 @@ package cache
 import (
 	"errors"
 	"fmt"
+	"io"
 	"sync"
 
 	"github.com/dgraph-io/badger/v2"
 	"github.com/dgrijalva/lfu-go"
-
 	"github.com/pyroscope-io/pyroscope/pkg/util/metrics"
+	"github.com/valyala/bytebufferpool"
 )
 
 type Cache struct {
@@ -94,17 +95,28 @@ func (cache *Cache) Put(key string, val interface{}) {
 	}
 }
 
+type serializable interface{ Serialize(io.Writer) error }
+
 func (cache *Cache) saveToDisk(key string, val interface{}) error {
-	// serialize the key and value
-	buf, err := cache.Bytes(key, val)
+	var buf []byte
+	var err error
+	if s, ok := val.(serializable); ok {
+		b := bytebufferpool.Get()
+		defer bytebufferpool.Put(b)
+		if err = s.Serialize(b); err == nil {
+			buf = b.Bytes()
+		}
+	} else {
+		// Note that `tree.Tree` does not satisfy serializable interface.
+		buf, err = cache.Bytes(key, val)
+	}
 	if err != nil {
 		return fmt.Errorf("serialize key and value: %v", err)
 	}
 
 	metrics.Count(cache.storageWriteCounter, 1)
-	// update the kv to badger
-	if err := cache.db.Update(func(txn *badger.Txn) error {
-		return txn.SetEntry(badger.NewEntry([]byte(cache.prefix+key), buf))
+	if err = cache.db.Update(func(txn *badger.Txn) error {
+		return txn.Set([]byte(cache.prefix+key), buf)
 	}); err != nil {
 		return fmt.Errorf("save to disk: %v", err)
 	}