Pebble pooled merger

claimtrie/change/change.go

@@ -78,9 +78,10 @@ func (c *Change) Marshal(enc *bytes.Buffer) error {
 		binary.BigEndian.PutUint32(temp[:4], uint32(len(c.SpentChildren)))
 		enc.Write(temp[:4])
 		for key := range c.SpentChildren {
-			binary.BigEndian.PutUint16(temp[:2], uint16(len(key))) // technically limited to 255; not sure we trust it
+			keySize := uint16(len(key))
+			binary.BigEndian.PutUint16(temp[:2], keySize) // technically limited to 255; not sure we trust it
 			enc.Write(temp[:2])
-			enc.WriteString(key)
+			enc.WriteString(key[:keySize])
 		}
 	} else {
 		binary.BigEndian.PutUint32(temp[:4], 0)

claimtrie/node/noderepo/pebble.go

@@ -2,7 +2,9 @@ package noderepo
 
 import (
 	"bytes"
+	"io"
 	"sort"
+	"sync"
 
 	"github.com/cockroachdb/pebble"
 	"github.com/lbryio/lbcd/claimtrie/change"
@@ -13,9 +15,76 @@ type Pebble struct {
 	db *pebble.DB
 }
 
+type pooledMerger struct {
+	values [][]byte
+	index  []int
+	pool   *sync.Pool
+	buffer []byte
+}
+
+func (a *pooledMerger) Len() int           { return len(a.index) }
+func (a *pooledMerger) Less(i, j int) bool { return a.index[i] < a.index[j] }
+func (a *pooledMerger) Swap(i, j int) {
+	a.index[i], a.index[j] = a.index[j], a.index[i]
+	a.values[i], a.values[j] = a.values[j], a.values[i]
+}
+
+func (a *pooledMerger) MergeNewer(value []byte) error {
+	vc := a.pool.Get().([]byte)[:0]
+	vc = append(vc, value...)
+	a.values = append(a.values, vc)
+	a.index = append(a.index, len(a.values))
+	return nil
+}
+
+func (a *pooledMerger) MergeOlder(value []byte) error {
+	vc := a.pool.Get().([]byte)[:0]
+	vc = append(vc, value...)
+	a.values = append(a.values, vc)
+	a.index = append(a.index, -len(a.values))
+	return nil
+}
+
+func (a *pooledMerger) Finish(includesBase bool) ([]byte, io.Closer, error) {
+	sort.Sort(a)
+
+	a.buffer = a.pool.Get().([]byte)[:0]
+	for i := range a.values {
+		a.buffer = append(a.buffer, a.values[i]...)
+	}
+
+	return a.buffer, a, nil
+}
+
+func (a *pooledMerger) Close() error {
+	for i := range a.values {
+		a.pool.Put(a.values[i])
+	}
+	a.pool.Put(a.buffer)
+	return nil
+}
+
 func NewPebble(path string) (*Pebble, error) {
 
-	db, err := pebble.Open(path, &pebble.Options{Cache: pebble.NewCache(64 << 20), BytesPerSync: 8 << 20, MaxOpenFiles: 2000})
+	mp := &sync.Pool{
+		New: func() interface{} {
+			return make([]byte, 0, 256)
+		},
+	}
+
+	db, err := pebble.Open(path, &pebble.Options{
+		Merger: &pebble.Merger{
+			Merge: func(key, value []byte) (pebble.ValueMerger, error) {
+				p := &pooledMerger{pool: mp}
+				return p, p.MergeNewer(value)
+			},
+			Name: pebble.DefaultMerger.Name, // yes, it's a lie
+		},
+		Cache:        pebble.NewCache(64 << 20),
+		BytesPerSync: 8 << 20,
+		MaxOpenFiles: 2000,
+	})
+
 	repo := &Pebble{db: db}
 
 	return repo, errors.Wrapf(err, "unable to open %s", path)