MerkleDB Improve Node Size Calculation (#1950)

Author: dboehm-avalabs

x/merkledb/db.go

@@ -1211,9 +1211,11 @@ func addPrefixToKey(bufferPool *sync.Pool, prefix []byte, key []byte) []byte {
 	return prefixedKey
 }
 
+// cacheEntrySize returns a rough approximation of the memory consumed by storing the path and node
 func cacheEntrySize(p path, n *node) int {
 	if n == nil {
 		return len(p)
 	}
-	return len(p) + len(n.bytes())
+	// nodes cache their bytes representation so the total memory consumed is roughly twice that
+	return len(p) + 2*len(n.bytes())
 }

x/merkledb/intermediate_node_db.go

@@ -27,7 +27,8 @@ type intermediateNodeDB struct {
 	// from the cache, which will call [OnEviction].
 	// A non-nil error returned from Put is considered fatal.
 	// Keys in [nodeCache] aren't prefixed with [intermediateNodePrefix].
-	nodeCache         onEvictCache[path, *node]
+	nodeCache onEvictCache[path, *node]
+	// the number of bytes to evict during an eviction batch
 	evictionBatchSize int
 	metrics           merkleMetrics
 }
@@ -57,6 +58,7 @@ func newIntermediateNodeDB(
 func (db *intermediateNodeDB) onEviction(key path, n *node) error {
 	writeBatch := db.baseDB.NewBatch()
 
+	totalSize := cacheEntrySize(key, n)
 	if err := db.addToBatch(writeBatch, key, n); err != nil {
 		_ = db.baseDB.Close()
 		return err
@@ -66,12 +68,14 @@ func (db *intermediateNodeDB) onEviction(key path, n *node) error {
 	// and write them to disk. We write a batch of them, rather than
 	// just [n], so that we don't immediately evict and write another
 	// node, because each time this method is called we do a disk write.
-	for writeBatch.Size() < db.evictionBatchSize {
+	// Evicts a total number of bytes, rather than a number of nodes
+	for totalSize < db.evictionBatchSize {
 		key, n, exists := db.nodeCache.removeOldest()
 		if !exists {
 			// The cache is empty.
 			break
 		}
+		totalSize += cacheEntrySize(key, n)
 		if err := db.addToBatch(writeBatch, key, n); err != nil {
 			_ = db.baseDB.Close()
 			return err

x/merkledb/intermediate_node_db_test.go

@@ -23,7 +23,8 @@ import (
 func TestIntermediateNodeDB(t *testing.T) {
 	require := require.New(t)
 
-	cacheSize := 100
+	// use exact multiple of node size so require.Equal(1, db.nodeCache.fifo.Len()) is correct later
+	cacheSize := 200
 	evictionBatchSize := cacheSize
 	baseDB := memdb.New()
 	db := newIntermediateNodeDB(
@@ -37,49 +38,40 @@ func TestIntermediateNodeDB(t *testing.T) {
 	)
 
 	// Put a key-node pair
-	key := newPath([]byte{0x01})
-	node1 := &node{
-		dbNode: dbNode{
-			value: maybe.Some([]byte{0x01}),
-		},
-	}
-	require.NoError(db.Put(key, node1))
+	node1Key := newPath([]byte{0x01})
+	node1 := newNode(nil, node1Key)
+	node1.setValue(maybe.Some([]byte{byte(0x01)}))
+	require.NoError(db.Put(node1Key, node1))
 
 	// Get the key-node pair from cache
-	node1Read, err := db.Get(key)
+	node1Read, err := db.Get(node1Key)
 	require.NoError(err)
 	require.Equal(node1, node1Read)
 
 	// Overwrite the key-node pair
-	node1Updated := &node{
-		dbNode: dbNode{
-			value: maybe.Some([]byte{0x02}),
-		},
-	}
-	require.NoError(db.Put(key, node1Updated))
+	node1Updated := newNode(nil, node1Key)
+	node1Updated.setValue(maybe.Some([]byte{byte(0x02)}))
+	require.NoError(db.Put(node1Key, node1Updated))
 
 	// Assert the key-node pair was overwritten
-	node1Read, err = db.Get(key)
+	node1Read, err = db.Get(node1Key)
 	require.NoError(err)
 	require.Equal(node1Updated, node1Read)
 
 	// Delete the key-node pair
-	require.NoError(db.Delete(key))
-	_, err = db.Get(key)
+	require.NoError(db.Delete(node1Key))
+	_, err = db.Get(node1Key)
 
 	// Assert the key-node pair was deleted
 	require.Equal(database.ErrNotFound, err)
 
 	// Put elements in the cache until it is full.
-	expectedSize := cacheEntrySize(key, nil)
+	expectedSize := 0
 	added := 0
 	for {
 		key := newPath([]byte{byte(added)})
-		node := &node{
-			dbNode: dbNode{
-				value: maybe.Some([]byte{byte(added)}),
-			},
-		}
+		node := newNode(nil, EmptyPath)
+		node.setValue(maybe.Some([]byte{byte(added)}))
 		newExpectedSize := expectedSize + cacheEntrySize(key, node)
 		if newExpectedSize > cacheSize {
 			// Don't trigger eviction.
@@ -97,24 +89,24 @@ func TestIntermediateNodeDB(t *testing.T) {
 	// Put one more element in the cache, which should trigger an eviction
 	// of all but 2 elements. 2 elements remain rather than 1 element because of
 	// the added key prefix increasing the size tracked by the batch.
-	key = newPath([]byte{byte(cacheSize)})
-	node := &node{
-		dbNode: dbNode{
-			value: maybe.Some([]byte{byte(cacheSize)}),
-		},
-	}
+	key := newPath([]byte{byte(added)})
+	node := newNode(nil, EmptyPath)
+	node.setValue(maybe.Some([]byte{byte(added)}))
 	require.NoError(db.Put(key, node))
 
 	// Assert cache has expected number of elements
-	require.Equal(2, db.nodeCache.fifo.Len())
+	require.Equal(1, db.nodeCache.fifo.Len())
 	gotKey, _, ok := db.nodeCache.fifo.Oldest()
 	require.True(ok)
-	require.Equal(newPath([]byte{byte(added - 1)}), gotKey)
+	require.Equal(newPath([]byte{byte(added)}), gotKey)
 
-	// Get a node from the base database (not cache)
-	nodeRead, err := db.Get(newPath([]byte{0x03}))
+	// Get a node from the base database
+	// Use an early key that has been evicted from the cache
+	_, inCache := db.nodeCache.Get(node1Key)
+	require.False(inCache)
+	nodeRead, err := db.Get(node1Key)
 	require.NoError(err)
-	require.Equal(maybe.Some([]byte{0x03}), nodeRead.value)
+	require.Equal(maybe.Some([]byte{0x01}), nodeRead.value)
 
 	// Flush the cache.
 	require.NoError(db.Flush())

x/merkledb/node.go

@@ -36,7 +36,7 @@ type child struct {
 	hasValue       bool
 }
 
-// node holds additional information on top of the dbNode that makes calulcations easier to do
+// node holds additional information on top of the dbNode that makes calculations easier to do
 type node struct {
 	dbNode
 	id          ids.ID
@@ -159,11 +159,9 @@ func (n *node) removeChild(child *node) {
 }
 
 // clone Returns a copy of [n].
-// nodeBytes is intentionally not included because it can cause a race.
-// nodes being evicted by the cache can write nodeBytes,
-// so reading them during the cloning would be a data race.
 // Note: value isn't cloned because it is never edited, only overwritten
 // if this ever changes, value will need to be copied as well
+// it is safe to clone all fields because they are only written/read while one or both of the db locks are held
 func (n *node) clone() *node {
 	return &node{
 		id:  n.id,
@@ -173,6 +171,7 @@ func (n *node) clone() *node {
 			children: maps.Clone(n.children),
 		},
 		valueDigest: n.valueDigest,
+		nodeBytes:   n.nodeBytes,
 	}
 }