triedb/pathdb: improve perf by separating nodes map (ethereum#31306)

This PR refactors the `nodeSet` structure in the path database to use
separate maps for account and storage trie nodes, resulting in
performance improvements. The change maintains the same API while
optimizing the internal data structure.

Author: weiihann

triedb/pathdb/nodes.go

@@ -36,8 +36,9 @@ import (
 // transition, typically corresponding to a block execution. It can also represent
 // the combined trie node set from several aggregated state transitions.
 type nodeSet struct {
-	size  uint64                                    // aggregated size of the trie node
-	nodes map[common.Hash]map[string]*trienode.Node // node set, mapped by owner and path
+	size         uint64                                    // aggregated size of the trie node
+	accountNodes map[string]*trienode.Node                 // account trie nodes, mapped by path
+	storageNodes map[common.Hash]map[string]*trienode.Node // storage trie nodes, mapped by owner and path
 }
 
 // newNodeSet constructs the set with the provided dirty trie nodes.
@@ -46,21 +47,30 @@ func newNodeSet(nodes map[common.Hash]map[string]*trienode.Node) *nodeSet {
 	if nodes == nil {
 		nodes = make(map[common.Hash]map[string]*trienode.Node)
 	}
-	s := &nodeSet{nodes: nodes}
+	s := &nodeSet{
+		accountNodes: make(map[string]*trienode.Node),
+		storageNodes: make(map[common.Hash]map[string]*trienode.Node),
+	}
+	for owner, subset := range nodes {
+		if owner == (common.Hash{}) {
+			s.accountNodes = subset
+		} else {
+			s.storageNodes[owner] = subset
+		}
+	}
 	s.computeSize()
 	return s
 }
 
 // computeSize calculates the database size of the held trie nodes.
 func (s *nodeSet) computeSize() {
 	var size uint64
-	for owner, subset := range s.nodes {
-		var prefix int
-		if owner != (common.Hash{}) {
-			prefix = common.HashLength // owner (32 bytes) for storage trie nodes
-		}
+	for path, n := range s.accountNodes {
+		size += uint64(len(n.Blob) + len(path))
+	}
+	for _, subset := range s.storageNodes {
 		for path, n := range subset {
-			size += uint64(prefix + len(n.Blob) + len(path))
+			size += uint64(common.HashLength + len(n.Blob) + len(path))
 		}
 	}
 	s.size = size
@@ -79,15 +89,18 @@ func (s *nodeSet) updateSize(delta int64) {
 
 // node retrieves the trie node with node path and its trie identifier.
 func (s *nodeSet) node(owner common.Hash, path []byte) (*trienode.Node, bool) {
-	subset, ok := s.nodes[owner]
-	if !ok {
-		return nil, false
+	// Account trie node
+	if owner == (common.Hash{}) {
+		n, ok := s.accountNodes[string(path)]
+		return n, ok
 	}
-	n, ok := subset[string(path)]
+	// Storage trie node
+	subset, ok := s.storageNodes[owner]
 	if !ok {
 		return nil, false
 	}
-	return n, true
+	n, ok := subset[string(path)]
+	return n, ok
 }
 
 // merge integrates the provided dirty nodes into the set. The provided nodeset
@@ -97,35 +110,44 @@ func (s *nodeSet) merge(set *nodeSet) {
 		delta     int64   // size difference resulting from node merging
 		overwrite counter // counter of nodes being overwritten
 	)
-	for owner, subset := range set.nodes {
-		var prefix int
-		if owner != (common.Hash{}) {
-			prefix = common.HashLength
+
+	// Merge account nodes
+	for path, n := range set.accountNodes {
+		if orig, exist := s.accountNodes[path]; !exist {
+			delta += int64(len(n.Blob) + len(path))
+		} else {
+			delta += int64(len(n.Blob) - len(orig.Blob))
+			overwrite.add(len(orig.Blob) + len(path))
 		}
-		current, exist := s.nodes[owner]
+		s.accountNodes[path] = n
+	}
+
+	// Merge storage nodes
+	for owner, subset := range set.storageNodes {
+		current, exist := s.storageNodes[owner]
 		if !exist {
 			for path, n := range subset {
-				delta += int64(prefix + len(n.Blob) + len(path))
+				delta += int64(common.HashLength + len(n.Blob) + len(path))
 			}
 			// Perform a shallow copy of the map for the subset instead of claiming it
 			// directly from the provided nodeset to avoid potential concurrent map
 			// read/write issues. The nodes belonging to the original diff layer remain
 			// accessible even after merging. Therefore, ownership of the nodes map
 			// should still belong to the original layer, and any modifications to it
 			// should be prevented.
-			s.nodes[owner] = maps.Clone(subset)
+			s.storageNodes[owner] = maps.Clone(subset)
 			continue
 		}
 		for path, n := range subset {
 			if orig, exist := current[path]; !exist {
-				delta += int64(prefix + len(n.Blob) + len(path))
+				delta += int64(common.HashLength + len(n.Blob) + len(path))
 			} else {
 				delta += int64(len(n.Blob) - len(orig.Blob))
-				overwrite.add(prefix + len(orig.Blob) + len(path))
+				overwrite.add(common.HashLength + len(orig.Blob) + len(path))
 			}
 			current[path] = n
 		}
-		s.nodes[owner] = current
+		s.storageNodes[owner] = current
 	}
 	overwrite.report(gcTrieNodeMeter, gcTrieNodeBytesMeter)
 	s.updateSize(delta)
@@ -136,34 +158,38 @@ func (s *nodeSet) merge(set *nodeSet) {
 func (s *nodeSet) revertTo(db ethdb.KeyValueReader, nodes map[common.Hash]map[string]*trienode.Node) {
 	var delta int64
 	for owner, subset := range nodes {
-		current, ok := s.nodes[owner]
-		if !ok {
-			panic(fmt.Sprintf("non-existent subset (%x)", owner))
-		}
-		for path, n := range subset {
-			orig, ok := current[path]
-			if !ok {
-				// There is a special case in merkle tree that one child is removed
-				// from a fullNode which only has two children, and then a new child
-				// with different position is immediately inserted into the fullNode.
-				// In this case, the clean child of the fullNode will also be marked
-				// as dirty because of node collapse and expansion. In case of database
-				// rollback, don't panic if this "clean" node occurs which is not
-				// present in buffer.
-				var blob []byte
-				if owner == (common.Hash{}) {
-					blob = rawdb.ReadAccountTrieNode(db, []byte(path))
-				} else {
-					blob = rawdb.ReadStorageTrieNode(db, owner, []byte(path))
+		if owner == (common.Hash{}) {
+			// Account trie nodes
+			for path, n := range subset {
+				orig, ok := s.accountNodes[path]
+				if !ok {
+					blob := rawdb.ReadAccountTrieNode(db, []byte(path))
+					if bytes.Equal(blob, n.Blob) {
+						continue
+					}
+					panic(fmt.Sprintf("non-existent account node (%v) blob: %v", path, crypto.Keccak256Hash(n.Blob).Hex()))
 				}
-				// Ignore the clean node in the case described above.
-				if bytes.Equal(blob, n.Blob) {
-					continue
+				s.accountNodes[path] = n
+				delta += int64(len(n.Blob)) - int64(len(orig.Blob))
+			}
+		} else {
+			// Storage trie nodes
+			current, ok := s.storageNodes[owner]
+			if !ok {
+				panic(fmt.Sprintf("non-existent subset (%x)", owner))
+			}
+			for path, n := range subset {
+				orig, ok := current[path]
+				if !ok {
+					blob := rawdb.ReadStorageTrieNode(db, owner, []byte(path))
+					if bytes.Equal(blob, n.Blob) {
+						continue
+					}
+					panic(fmt.Sprintf("non-existent storage node (%x %v) blob: %v", owner, path, crypto.Keccak256Hash(n.Blob).Hex()))
 				}
-				panic(fmt.Sprintf("non-existent node (%x %v) blob: %v", owner, path, crypto.Keccak256Hash(n.Blob).Hex()))
+				current[path] = n
+				delta += int64(len(n.Blob)) - int64(len(orig.Blob))
 			}
-			current[path] = n
-			delta += int64(len(n.Blob)) - int64(len(orig.Blob))
 		}
 	}
 	s.updateSize(delta)
@@ -184,8 +210,21 @@ type journalNodes struct {
 
 // encode serializes the content of trie nodes into the provided writer.
 func (s *nodeSet) encode(w io.Writer) error {
-	nodes := make([]journalNodes, 0, len(s.nodes))
-	for owner, subset := range s.nodes {
+	nodes := make([]journalNodes, 0, len(s.storageNodes)+1)
+
+	// Encode account nodes
+	if len(s.accountNodes) > 0 {
+		entry := journalNodes{Owner: common.Hash{}}
+		for path, node := range s.accountNodes {
+			entry.Nodes = append(entry.Nodes, journalNode{
+				Path: []byte(path),
+				Blob: node.Blob,
+			})
+		}
+		nodes = append(nodes, entry)
+	}
+	// Encode storage nodes
+	for owner, subset := range s.storageNodes {
 		entry := journalNodes{Owner: owner}
 		for path, node := range subset {
 			entry.Nodes = append(entry.Nodes, journalNode{
@@ -204,43 +243,61 @@ func (s *nodeSet) decode(r *rlp.Stream) error {
 	if err := r.Decode(&encoded); err != nil {
 		return fmt.Errorf("load nodes: %v", err)
 	}
-	nodes := make(map[common.Hash]map[string]*trienode.Node)
+	s.accountNodes = make(map[string]*trienode.Node)
+	s.storageNodes = make(map[common.Hash]map[string]*trienode.Node)
+
 	for _, entry := range encoded {
-		subset := make(map[string]*trienode.Node)
-		for _, n := range entry.Nodes {
-			if len(n.Blob) > 0 {
-				subset[string(n.Path)] = trienode.New(crypto.Keccak256Hash(n.Blob), n.Blob)
-			} else {
-				subset[string(n.Path)] = trienode.NewDeleted()
+		if entry.Owner == (common.Hash{}) {
+			// Account nodes
+			for _, n := range entry.Nodes {
+				if len(n.Blob) > 0 {
+					s.accountNodes[string(n.Path)] = trienode.New(crypto.Keccak256Hash(n.Blob), n.Blob)
+				} else {
+					s.accountNodes[string(n.Path)] = trienode.NewDeleted()
+				}
+			}
+		} else {
+			// Storage nodes
+			subset := make(map[string]*trienode.Node)
+			for _, n := range entry.Nodes {
+				if len(n.Blob) > 0 {
+					subset[string(n.Path)] = trienode.New(crypto.Keccak256Hash(n.Blob), n.Blob)
+				} else {
+					subset[string(n.Path)] = trienode.NewDeleted()
+				}
 			}
+			s.storageNodes[entry.Owner] = subset
 		}
-		nodes[entry.Owner] = subset
 	}
-	s.nodes = nodes
 	s.computeSize()
 	return nil
 }
 
 // write flushes nodes into the provided database batch as a whole.
 func (s *nodeSet) write(batch ethdb.Batch, clean *fastcache.Cache) int {
-	return writeNodes(batch, s.nodes, clean)
+	nodes := make(map[common.Hash]map[string]*trienode.Node)
+	if len(s.accountNodes) > 0 {
+		nodes[common.Hash{}] = s.accountNodes
+	}
+	for owner, subset := range s.storageNodes {
+		nodes[owner] = subset
+	}
+	return writeNodes(batch, nodes, clean)
 }
 
 // reset clears all cached trie node data.
 func (s *nodeSet) reset() {
-	s.nodes = make(map[common.Hash]map[string]*trienode.Node)
+	s.accountNodes = make(map[string]*trienode.Node)
+	s.storageNodes = make(map[common.Hash]map[string]*trienode.Node)
 	s.size = 0
 }
 
 // dbsize returns the approximate size of db write.
 func (s *nodeSet) dbsize() int {
 	var m int
-	for owner, nodes := range s.nodes {
-		if owner == (common.Hash{}) {
-			m += len(nodes) * len(rawdb.TrieNodeAccountPrefix) // database key prefix
-		} else {
-			m += len(nodes) * (len(rawdb.TrieNodeStoragePrefix)) // database key prefix
-		}
+	m += len(s.accountNodes) * len(rawdb.TrieNodeAccountPrefix) // database key prefix
+	for _, nodes := range s.storageNodes {
+		m += len(nodes) * (len(rawdb.TrieNodeStoragePrefix)) // database key prefix
 	}
 	return m + int(s.size)
 }