Fix intermediate node caching (#2585)

Co-authored-by: Dan Laine <[email protected]>
ava-labs · Jan 4, 2024 · 7c82a5b · 7c82a5b
1 parent 72dc442
commit 7c82a5b
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 36 deletions.
diff --git a/x/merkledb/db.go b/x/merkledb/db.go
@@ -259,23 +259,22 @@ func newDatabase(
 			return make([]byte, 0, defaultBufferLength)
 		},
 	}
-	iNodeDB, err := newIntermediateNodeDB(
-		db,
-		bufferPool,
-		metrics,
-		int(config.IntermediateNodeCacheSize),
-		int(config.IntermediateWriteBufferSize),
-		int(config.IntermediateWriteBatchSize),
-		BranchFactorToTokenSize[config.BranchFactor])
-	if err != nil {
-		return nil, err
-	}
 
 	trieDB := &merkleDB{
-		metrics:              metrics,
-		baseDB:               db,
-		intermediateNodeDB:   iNodeDB,
-		valueNodeDB:          newValueNodeDB(db, bufferPool, metrics, int(config.ValueNodeCacheSize)),
+		metrics: metrics,
+		baseDB:  db,
+		intermediateNodeDB: newIntermediateNodeDB(
+			db,
+			bufferPool,
+			metrics,
+			int(config.IntermediateNodeCacheSize),
+			int(config.IntermediateWriteBufferSize),
+			int(config.IntermediateWriteBatchSize),
+			BranchFactorToTokenSize[config.BranchFactor]),
+		valueNodeDB: newValueNodeDB(db,
+			bufferPool,
+			metrics,
+			int(config.ValueNodeCacheSize)),
 		history:              newTrieHistory(int(config.HistoryLength)),
 		debugTracer:          getTracerIfEnabled(config.TraceLevel, DebugTrace, config.Tracer),
 		infoTracer:           getTracerIfEnabled(config.TraceLevel, InfoTrace, config.Tracer),

diff --git a/x/merkledb/intermediate_node_db.go b/x/merkledb/intermediate_node_db.go
@@ -4,7 +4,6 @@
 package merkledb
 
 import (
-	"errors"
 	"sync"
 
 	"github.com/ava-labs/avalanchego/cache"
@@ -14,8 +13,6 @@ import (
 
 const defaultBufferLength = 256
 
-var errCacheSizeTooSmall = errors.New("cache size must be larger than or equal to write buffer size")
-
 // Holds intermediate nodes. That is, those without values.
 // Changes to this database aren't written to [baseDB] until
 // they're evicted from the [nodeCache] or Flush is called.
@@ -51,10 +48,7 @@ func newIntermediateNodeDB(
 	writeBufferSize int,
 	evictionBatchSize int,
 	tokenSize int,
-) (*intermediateNodeDB, error) {
-	if cacheSize < writeBufferSize {
-		return nil, errCacheSizeTooSmall
-	}
+) *intermediateNodeDB {
 	result := &intermediateNodeDB{
 		metrics:           metrics,
 		baseDB:            db,
@@ -69,7 +63,7 @@ func newIntermediateNodeDB(
 		result.onEviction,
 	)
 
-	return result, nil
+	return result
 }
 
 // A non-nil error is considered fatal and closes [db.baseDB].
@@ -125,6 +119,15 @@ func (db *intermediateNodeDB) Get(key Key) (*node, error) {
 	}
 	db.metrics.IntermediateNodeCacheMiss()
 
+	if cachedValue, isCached := db.writeBuffer.Get(key); isCached {
+		db.metrics.IntermediateNodeCacheHit()
+		if cachedValue == nil {
+			return nil, database.ErrNotFound
+		}
+		return cachedValue, nil
+	}
+	db.metrics.IntermediateNodeCacheMiss()
+
 	dbKey := db.constructDBKey(key)
 	db.metrics.DatabaseNodeRead()
 	nodeBytes, err := db.baseDB.Get(dbKey)

diff --git a/x/merkledb/intermediate_node_db_test.go b/x/merkledb/intermediate_node_db_test.go
@@ -33,7 +33,7 @@ func Test_IntermediateNodeDB(t *testing.T) {
 
 	evictionBatchSize := bufferSize
 	baseDB := memdb.New()
-	db, err := newIntermediateNodeDB(
+	db := newIntermediateNodeDB(
 		baseDB,
 		&sync.Pool{
 			New: func() interface{} { return make([]byte, 0) },
@@ -44,7 +44,6 @@ func Test_IntermediateNodeDB(t *testing.T) {
 		evictionBatchSize,
 		4,
 	)
-	require.NoError(err)
 
 	// Put a key-node pair
 	node1Key := ToKey([]byte{0x01})
@@ -148,7 +147,7 @@ func FuzzIntermediateNodeDBConstructDBKey(f *testing.F) {
 	) {
 		require := require.New(t)
 		for _, tokenSize := range validTokenSizes {
-			db, err := newIntermediateNodeDB(
+			db := newIntermediateNodeDB(
 				baseDB,
 				&sync.Pool{
 					New: func() interface{} { return make([]byte, 0) },
@@ -159,7 +158,6 @@ func FuzzIntermediateNodeDBConstructDBKey(f *testing.F) {
 				evictionBatchSize,
 				tokenSize,
 			)
-			require.NoError(err)
 
 			p := ToKey(key)
 			uBitLength := tokenLength * uint(tokenSize)
@@ -192,7 +190,7 @@ func Test_IntermediateNodeDB_ConstructDBKey_DirtyBuffer(t *testing.T) {
 	bufferSize := 200
 	evictionBatchSize := bufferSize
 	baseDB := memdb.New()
-	db, err := newIntermediateNodeDB(
+	db := newIntermediateNodeDB(
 		baseDB,
 		&sync.Pool{
 			New: func() interface{} { return make([]byte, 0) },
@@ -204,8 +202,6 @@ func Test_IntermediateNodeDB_ConstructDBKey_DirtyBuffer(t *testing.T) {
 		4,
 	)
 
-	require.NoError(err)
-
 	db.bufferPool.Put([]byte{0xFF, 0xFF, 0xFF})
 	constructedKey := db.constructDBKey(ToKey([]byte{}))
 	require.Len(constructedKey, 2)
@@ -231,7 +227,7 @@ func TestIntermediateNodeDBClear(t *testing.T) {
 	bufferSize := 200
 	evictionBatchSize := bufferSize
 	baseDB := memdb.New()
-	db, err := newIntermediateNodeDB(
+	db := newIntermediateNodeDB(
 		baseDB,
 		&sync.Pool{
 			New: func() interface{} { return make([]byte, 0) },
@@ -243,8 +239,6 @@ func TestIntermediateNodeDBClear(t *testing.T) {
 		4,
 	)
 
-	require.NoError(err)
-
 	for _, b := range [][]byte{{1}, {2}, {3}} {
 		require.NoError(db.Put(ToKey(b), newNode(ToKey(b))))
 	}
@@ -269,7 +263,7 @@ func TestIntermediateNodeDBDeleteEmptyKey(t *testing.T) {
 	bufferSize := 200
 	evictionBatchSize := bufferSize
 	baseDB := memdb.New()
-	db, err := newIntermediateNodeDB(
+	db := newIntermediateNodeDB(
 		baseDB,
 		&sync.Pool{
 			New: func() interface{} { return make([]byte, 0) },
@@ -281,8 +275,6 @@ func TestIntermediateNodeDBDeleteEmptyKey(t *testing.T) {
 		4,
 	)
 
-	require.NoError(err)
-
 	emptyKey := ToKey([]byte{})
 	require.NoError(db.Put(emptyKey, newNode(emptyKey)))
 	require.NoError(db.Flush())