Reuse key buffers during hashing

x/merkledb/bytes_pool.go

@@ -0,0 +1,67 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package merkledb
+
+import "sync"
+
+type bytesPool struct {
+	slots     chan struct{}
+	bytesLock sync.Mutex
+	bytes     [][]byte
+}
+
+func newBytesPool(numSlots int) *bytesPool {
+	return &bytesPool{
+		slots: make(chan struct{}, numSlots),
+		bytes: make([][]byte, 0, numSlots),
+	}
+}
+
+func (p *bytesPool) Acquire() []byte {
+	p.slots <- struct{}{}
+	return p.pop()
+}
+
+func (p *bytesPool) TryAcquire() ([]byte, bool) {
+	select {
+	case p.slots <- struct{}{}:
+		return p.pop(), true
+	default:
+		return nil, false
+	}
+}
+
+func (p *bytesPool) pop() []byte {
+	p.bytesLock.Lock()
+	defer p.bytesLock.Unlock()
+
+	numBytes := len(p.bytes)
+	if numBytes == 0 {
+		return nil
+	}
+
+	b := p.bytes[numBytes-1]
+	p.bytes = p.bytes[:numBytes-1]
+	return b
+}
+
+func (p *bytesPool) Release(b []byte) {
+	select {
+	case <-p.slots:
+	default:
+		panic("release of unacquired semaphore")
+	}
+
+	p.bytesLock.Lock()
+	defer p.bytesLock.Unlock()
+
+	p.bytes = append(p.bytes, b)
+}
+
+func setLength(b []byte, size int) []byte {
+	if size <= cap(b) {
+		return b[:size]
+	}
+	return append(b[:cap(b)], make([]byte, size-cap(b))...)
+}

x/merkledb/bytes_pool_test.go

@@ -0,0 +1,46 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package merkledb
+
+import "testing"
+
+func Benchmark_BytesPool_Acquire(b *testing.B) {
+	s := newBytesPool(b.N)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		s.Acquire()
+	}
+}
+
+func Benchmark_BytesPool_Release(b *testing.B) {
+	s := newBytesPool(b.N)
+	for i := 0; i < b.N; i++ {
+		s.Acquire()
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		s.Release(nil)
+	}
+}
+
+func Benchmark_BytesPool_TryAcquire_Success(b *testing.B) {
+	s := newBytesPool(b.N)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		s.TryAcquire()
+	}
+}
+
+func Benchmark_BytesPool_TryAcquire_Failure(b *testing.B) {
+	s := newBytesPool(1)
+	s.Acquire()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		s.TryAcquire()
+	}
+}

x/merkledb/db.go

@@ -15,7 +15,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/otel/attribute"
 	"golang.org/x/exp/maps"
-	"golang.org/x/sync/semaphore"
 
 	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/ids"
@@ -216,9 +215,10 @@ type merkleDB struct {
 	// Valid children of this trie.
 	childViews []*view
 
-	// calculateNodeIDsSema controls the number of goroutines inside
-	// [calculateNodeIDsHelper] at any given time.
-	calculateNodeIDsSema *semaphore.Weighted
+	// hashNodesKeyPool controls the number of goroutines that are created
+	// inside [hashChangedNode] at any given time and provides slices for the
+	// keys needed while hashing.
+	hashNodesKeyPool *bytesPool
 
 	tokenSize int
 }
@@ -242,9 +242,9 @@ func newDatabase(
 		return nil, err
 	}
 
-	rootGenConcurrency := uint(runtime.NumCPU())
+	rootGenConcurrency := runtime.NumCPU()
 	if config.RootGenConcurrency != 0 {
-		rootGenConcurrency = config.RootGenConcurrency
+		rootGenConcurrency = int(config.RootGenConcurrency)
 	}
 
 	// Share a sync.Pool of []byte between the intermediateNodeDB and valueNodeDB
@@ -270,12 +270,12 @@ func newDatabase(
 			bufferPool,
 			metrics,
 			int(config.ValueNodeCacheSize)),
-		history:              newTrieHistory(int(config.HistoryLength)),
-		debugTracer:          getTracerIfEnabled(config.TraceLevel, DebugTrace, config.Tracer),
-		infoTracer:           getTracerIfEnabled(config.TraceLevel, InfoTrace, config.Tracer),
-		childViews:           make([]*view, 0, defaultPreallocationSize),
-		calculateNodeIDsSema: semaphore.NewWeighted(int64(rootGenConcurrency)),
-		tokenSize:            BranchFactorToTokenSize[config.BranchFactor],
+		history:          newTrieHistory(int(config.HistoryLength)),
+		debugTracer:      getTracerIfEnabled(config.TraceLevel, DebugTrace, config.Tracer),
+		infoTracer:       getTracerIfEnabled(config.TraceLevel, InfoTrace, config.Tracer),
+		childViews:       make([]*view, 0, defaultPreallocationSize),
+		hashNodesKeyPool: newBytesPool(rootGenConcurrency),
+		tokenSize:        BranchFactorToTokenSize[config.BranchFactor],
 	}
 
 	if err := trieDB.initializeRoot(); err != nil {

x/merkledb/history.go

@@ -57,7 +57,7 @@ type changeSummary struct {
 	// The ID of the trie after these changes.
 	rootID ids.ID
 	// The root before/after this change.
-	// Set in [calculateNodeIDs].
+	// Set in [applyValueChanges].
 	rootChange change[maybe.Maybe[*node]]
 	nodes      map[Key]*change[*node]
 	values     map[Key]*change[maybe.Maybe[[]byte]]

x/merkledb/trie_test.go

@@ -21,7 +21,7 @@ import (
 func getNodeValue(t Trie, key string) ([]byte, error) {
 	path := ToKey([]byte(key))
 	if asView, ok := t.(*view); ok {
-		if err := asView.calculateNodeIDs(context.Background()); err != nil {
+		if err := asView.applyValueChanges(context.Background()); err != nil {
 			return nil, err
 		}
 	}
@@ -131,7 +131,7 @@ func TestVisitPathToKey(t *testing.T) {
 	require.NoError(err)
 	require.IsType(&view{}, trieIntf)
 	trie = trieIntf.(*view)
-	require.NoError(trie.calculateNodeIDs(context.Background()))
+	require.NoError(trie.applyValueChanges(context.Background()))
 
 	nodePath = make([]*node, 0, 1)
 	require.NoError(visitPathToKey(trie, ToKey(key1), func(n *node) error {
@@ -156,7 +156,7 @@ func TestVisitPathToKey(t *testing.T) {
 	require.NoError(err)
 	require.IsType(&view{}, trieIntf)
 	trie = trieIntf.(*view)
-	require.NoError(trie.calculateNodeIDs(context.Background()))
+	require.NoError(trie.applyValueChanges(context.Background()))
 
 	nodePath = make([]*node, 0, 2)
 	require.NoError(visitPathToKey(trie, ToKey(key2), func(n *node) error {
@@ -185,7 +185,7 @@ func TestVisitPathToKey(t *testing.T) {
 	require.NoError(err)
 	require.IsType(&view{}, trieIntf)
 	trie = trieIntf.(*view)
-	require.NoError(trie.calculateNodeIDs(context.Background()))
+	require.NoError(trie.applyValueChanges(context.Background()))
 
 	// Trie is:
 	//    []
@@ -775,7 +775,7 @@ func Test_Trie_ChainDeletion(t *testing.T) {
 	)
 	require.NoError(err)
 
-	require.NoError(newTrie.(*view).calculateNodeIDs(context.Background()))
+	require.NoError(newTrie.(*view).applyValueChanges(context.Background()))
 	maybeRoot := newTrie.getRoot()
 	require.NoError(err)
 	require.True(maybeRoot.HasValue())
@@ -794,7 +794,7 @@ func Test_Trie_ChainDeletion(t *testing.T) {
 		},
 	)
 	require.NoError(err)
-	require.NoError(newTrie.(*view).calculateNodeIDs(context.Background()))
+	require.NoError(newTrie.(*view).applyValueChanges(context.Background()))
 
 	// trie should be empty
 	root := newTrie.getRoot()
@@ -861,7 +861,7 @@ func Test_Trie_NodeCollapse(t *testing.T) {
 	)
 	require.NoError(err)
 
-	require.NoError(trie.(*view).calculateNodeIDs(context.Background()))
+	require.NoError(trie.(*view).applyValueChanges(context.Background()))
 
 	for _, kv := range kvs {
 		node, err := trie.getEditableNode(ToKey(kv.Key), true)
@@ -888,7 +888,7 @@ func Test_Trie_NodeCollapse(t *testing.T) {
 	)
 	require.NoError(err)
 
-	require.NoError(trie.(*view).calculateNodeIDs(context.Background()))
+	require.NoError(trie.(*view).applyValueChanges(context.Background()))
 
 	for _, kv := range deletedKVs {
 		_, err := trie.getEditableNode(ToKey(kv.Key), true)