Reduce allocations on insert and remove (#2201)

Signed-off-by: David Boehm <[email protected]> Signed-off-by: Dan Laine <[email protected]> Co-authored-by: Darioush Jalali <[email protected]> Co-authored-by: Dan Laine <[email protected]>
ava-labs · Oct 26, 2023 · e933587 · e933587
1 parent 638000c
commit e933587
Show file tree

Hide file tree

Showing 3 changed files with 135 additions and 178 deletions.
diff --git a/x/merkledb/db.go b/x/merkledb/db.go
@@ -474,19 +474,14 @@ func (db *merkleDB) PrefetchPath(key []byte) error {
 }
 
 func (db *merkleDB) prefetchPath(view *trieView, keyBytes []byte) error {
-	pathToKey, err := view.getPathTo(db.toKey(keyBytes))
-	if err != nil {
-		return err
-	}
-	for _, n := range pathToKey {
-		if n.hasValue() {
-			db.valueNodeDB.nodeCache.Put(n.key, n)
-		} else if err := db.intermediateNodeDB.nodeCache.Put(n.key, n); err != nil {
-			return err
+	return view.visitPathToKey(db.toKey(keyBytes), func(n *node) error {
+		if !n.hasValue() {
+			return db.intermediateNodeDB.nodeCache.Put(n.key, n)
 		}
-	}
 
-	return nil
+		db.valueNodeDB.nodeCache.Put(n.key, n)
+		return nil
+	})
 }
 
 func (db *merkleDB) Get(key []byte) ([]byte, error) {

diff --git a/x/merkledb/trie_test.go b/x/merkledb/trie_test.go
@@ -23,40 +23,35 @@ func getNodeValue(t ReadOnlyTrie, key string) ([]byte, error) {
 }
 
 func getNodeValueWithBranchFactor(t ReadOnlyTrie, key string, bf BranchFactor) ([]byte, error) {
+	var view *trieView
 	if asTrieView, ok := t.(*trieView); ok {
 		if err := asTrieView.calculateNodeIDs(context.Background()); err != nil {
 			return nil, err
 		}
-		path := ToKey([]byte(key), bf)
-		nodePath, err := asTrieView.getPathTo(path)
-		if err != nil {
-			return nil, err
-		}
-		closestNode := nodePath[len(nodePath)-1]
-		if closestNode.key != path || closestNode == nil {
-			return nil, database.ErrNotFound
-		}
-
-		return closestNode.value.Value(), nil
+		view = asTrieView
 	}
 	if asDatabases, ok := t.(*merkleDB); ok {
-		view, err := asDatabases.NewView(context.Background(), ViewChanges{})
+		dbView, err := asDatabases.NewView(context.Background(), ViewChanges{})
 		if err != nil {
 			return nil, err
 		}
-		path := ToKey([]byte(key), bf)
-		nodePath, err := view.(*trieView).getPathTo(path)
-		if err != nil {
-			return nil, err
-		}
-		closestNode := nodePath[len(nodePath)-1]
-		if closestNode.key != path || closestNode == nil {
-			return nil, database.ErrNotFound
-		}
+		view = dbView.(*trieView)
+	}
 
-		return closestNode.value.Value(), nil
+	path := ToKey([]byte(key), bf)
+	var result *node
+	err := view.visitPathToKey(path, func(n *node) error {
+		result = n
+		return nil
+	})
+	if err != nil {
+		return nil, err
 	}
-	return nil, nil
+	if result.key != path || result == nil {
+		return nil, database.ErrNotFound
+	}
+
+	return result.value.Value(), nil
 }
 
 func Test_GetValue_Safety(t *testing.T) {
@@ -116,7 +111,7 @@ func Test_GetValues_Safety(t *testing.T) {
 	require.Equal([]byte{0}, trieVals[0])
 }
 
-func TestTrieViewGetPathTo(t *testing.T) {
+func TestTrieViewVisitPathToKey(t *testing.T) {
 	require := require.New(t)
 
 	db, err := getBasicDB()
@@ -127,8 +122,11 @@ func TestTrieViewGetPathTo(t *testing.T) {
 	require.IsType(&trieView{}, trieIntf)
 	trie := trieIntf.(*trieView)
 
-	nodePath, err := trie.getPathTo(ToKey(nil, BranchFactor16))
-	require.NoError(err)
+	var nodePath []*node
+	require.NoError(trie.visitPathToKey(ToKey(nil, BranchFactor16), func(n *node) error {
+		nodePath = append(nodePath, n)
+		return nil
+	}))
 
 	// Just the root
 	require.Len(nodePath, 1)
@@ -149,8 +147,11 @@ func TestTrieViewGetPathTo(t *testing.T) {
 	trie = trieIntf.(*trieView)
 	require.NoError(trie.calculateNodeIDs(context.Background()))
 
-	nodePath, err = trie.getPathTo(ToKey(key1, BranchFactor16))
-	require.NoError(err)
+	nodePath = make([]*node, 0, 2)
+	require.NoError(trie.visitPathToKey(ToKey(key1, BranchFactor16), func(n *node) error {
+		nodePath = append(nodePath, n)
+		return nil
+	}))
 
 	// Root and 1 value
 	require.Len(nodePath, 2)
@@ -172,8 +173,11 @@ func TestTrieViewGetPathTo(t *testing.T) {
 	trie = trieIntf.(*trieView)
 	require.NoError(trie.calculateNodeIDs(context.Background()))
 
-	nodePath, err = trie.getPathTo(ToKey(key2, BranchFactor16))
-	require.NoError(err)
+	nodePath = make([]*node, 0, 3)
+	require.NoError(trie.visitPathToKey(ToKey(key2, BranchFactor16), func(n *node) error {
+		nodePath = append(nodePath, n)
+		return nil
+	}))
 	require.Len(nodePath, 3)
 	require.Equal(trie.root, nodePath[0])
 	require.Equal(ToKey(key1, BranchFactor16), nodePath[1].key)
@@ -194,33 +198,45 @@ func TestTrieViewGetPathTo(t *testing.T) {
 	trie = trieIntf.(*trieView)
 	require.NoError(trie.calculateNodeIDs(context.Background()))
 
-	nodePath, err = trie.getPathTo(ToKey(key3, BranchFactor16))
-	require.NoError(err)
+	nodePath = make([]*node, 0, 2)
+	require.NoError(trie.visitPathToKey(ToKey(key3, BranchFactor16), func(n *node) error {
+		nodePath = append(nodePath, n)
+		return nil
+	}))
 	require.Len(nodePath, 2)
 	require.Equal(trie.root, nodePath[0])
 	require.Equal(ToKey(key3, BranchFactor16), nodePath[1].key)
 
 	// Other key path not affected
-	nodePath, err = trie.getPathTo(ToKey(key2, BranchFactor16))
-	require.NoError(err)
+	nodePath = make([]*node, 0, 3)
+	require.NoError(trie.visitPathToKey(ToKey(key2, BranchFactor16), func(n *node) error {
+		nodePath = append(nodePath, n)
+		return nil
+	}))
 	require.Len(nodePath, 3)
 	require.Equal(trie.root, nodePath[0])
 	require.Equal(ToKey(key1, BranchFactor16), nodePath[1].key)
 	require.Equal(ToKey(key2, BranchFactor16), nodePath[2].key)
 
 	// Gets closest node when key doesn't exist
 	key4 := []byte{0, 1, 2}
-	nodePath, err = trie.getPathTo(ToKey(key4, BranchFactor16))
-	require.NoError(err)
+	nodePath = make([]*node, 0, 3)
+	require.NoError(trie.visitPathToKey(ToKey(key4, BranchFactor16), func(n *node) error {
+		nodePath = append(nodePath, n)
+		return nil
+	}))
 	require.Len(nodePath, 3)
 	require.Equal(trie.root, nodePath[0])
 	require.Equal(ToKey(key1, BranchFactor16), nodePath[1].key)
 	require.Equal(ToKey(key2, BranchFactor16), nodePath[2].key)
 
 	// Gets just root when key doesn't exist and no key shares a prefix
 	key5 := []byte{128}
-	nodePath, err = trie.getPathTo(ToKey(key5, BranchFactor16))
-	require.NoError(err)
+	nodePath = make([]*node, 0, 1)
+	require.NoError(trie.visitPathToKey(ToKey(key5, BranchFactor16), func(n *node) error {
+		nodePath = append(nodePath, n)
+		return nil
+	}))
 	require.Len(nodePath, 1)
 	require.Equal(trie.root, nodePath[0])
 }