triedb/pathdb: introduce lookup structure to optimize state access (#30971)

This pull request introduces a mechanism to improve state lookup
efficiency in pathdb by maintaining a lookup structure that eliminates
unnecessary iteration over diff layers.

The core idea is to track a mutation history for each dirty state entry
residing in the diff layers. This history records the state roots of all layers
in which the entry was modified, sorted from oldest to newest.

During state lookup, this mutation history is queried to find the most
recent layer whose state root either matches the target root or is a
descendant of it. This allows us to quickly identify the layer containing
the relevant data, avoiding the need to iterate through all diff layers from
top to bottom.

Besides, the overhead for state lookup is constant, no matter how many
diff layers are retained in the pathdb, which unlocks the potential to hold
more diff layers.

Of course, maintaining this lookup structure introduces some overhead.
For each state transition, we need to:
(a) update the mutation records for the modified state entries, and
(b) remove stale mutation records associated with outdated layers.

On our benchmark machine, it will introduce around 1ms overhead which is
acceptable.

Author: rjl493456442

triedb/pathdb/database.go

@@ -487,7 +487,7 @@ func (db *Database) Enable(root common.Hash) error {
 
 	// Re-construct a new disk layer backed by persistent state
 	// and schedule the state snapshot generation if it's permitted.
-	db.tree.reset(generateSnapshot(db, root, db.isVerkle || db.config.SnapshotNoBuild))
+	db.tree.init(generateSnapshot(db, root, db.isVerkle || db.config.SnapshotNoBuild))
 	log.Info("Rebuilt trie database", "root", root)
 	return nil
 }
@@ -529,7 +529,7 @@ func (db *Database) Recover(root common.Hash) error {
 		// reset layer with newly created disk layer. It must be
 		// done after each revert operation, otherwise the new
 		// disk layer won't be accessible from outside.
-		db.tree.reset(dl)
+		db.tree.init(dl)
 	}
 	rawdb.DeleteTrieJournal(db.diskdb)
 

triedb/pathdb/difflayer.go

@@ -156,7 +156,7 @@ func (dl *diffLayer) update(root common.Hash, id uint64, block uint64, nodes *no
 }
 
 // persist flushes the diff layer and all its parent layers to disk layer.
-func (dl *diffLayer) persist(force bool) (layer, error) {
+func (dl *diffLayer) persist(force bool) (*diskLayer, error) {
 	if parent, ok := dl.parentLayer().(*diffLayer); ok {
 		// Hold the lock to prevent any read operation until the new
 		// parent is linked correctly.
@@ -183,7 +183,7 @@ func (dl *diffLayer) size() uint64 {
 
 // diffToDisk merges a bottom-most diff into the persistent disk layer underneath
 // it. The method will panic if called onto a non-bottom-most diff layer.
-func diffToDisk(layer *diffLayer, force bool) (layer, error) {
+func diffToDisk(layer *diffLayer, force bool) (*diskLayer, error) {
 	disk, ok := layer.parentLayer().(*diskLayer)
 	if !ok {
 		panic(fmt.Sprintf("unknown layer type: %T", layer.parentLayer()))

triedb/pathdb/disklayer.go

@@ -94,15 +94,6 @@ func (dl *diskLayer) setGenerator(generator *generator) {
 	dl.generator = generator
 }
 
-// isStale return whether this layer has become stale (was flattened across) or if
-// it's still live.
-func (dl *diskLayer) isStale() bool {
-	dl.lock.RLock()
-	defer dl.lock.RUnlock()
-
-	return dl.stale
-}
-
 // markStale sets the stale flag as true.
 func (dl *diskLayer) markStale() {
 	dl.lock.Lock()

triedb/pathdb/layertree.go

@@ -31,29 +31,50 @@ import (
 // thread-safe to use. However, callers need to ensure the thread-safety
 // of the referenced layer by themselves.
 type layerTree struct {
-	lock   sync.RWMutex
+	base   *diskLayer
 	layers map[common.Hash]layer
+
+	// descendants is a two-dimensional map where the keys represent
+	// an ancestor state root, and the values are the state roots of
+	// all its descendants.
+	//
+	// For example: r -> [c1, c2, ..., cn], where c1 through cn are
+	// the descendants of state r.
+	//
+	// This map includes all the existing diff layers and the disk layer.
+	descendants map[common.Hash]map[common.Hash]struct{}
+	lookup      *lookup
+	lock        sync.RWMutex
 }
 
 // newLayerTree constructs the layerTree with the given head layer.
 func newLayerTree(head layer) *layerTree {
 	tree := new(layerTree)
-	tree.reset(head)
+	tree.init(head)
 	return tree
 }
 
-// reset initializes the layerTree by the given head layer.
-// All the ancestors will be iterated out and linked in the tree.
-func (tree *layerTree) reset(head layer) {
+// init initializes the layerTree by the given head layer.
+func (tree *layerTree) init(head layer) {
 	tree.lock.Lock()
 	defer tree.lock.Unlock()
 
-	var layers = make(map[common.Hash]layer)
-	for head != nil {
-		layers[head.rootHash()] = head
-		head = head.parentLayer()
+	current := head
+	tree.layers = make(map[common.Hash]layer)
+	tree.descendants = make(map[common.Hash]map[common.Hash]struct{})
+
+	for {
+		tree.layers[current.rootHash()] = current
+		tree.fillAncestors(current)
+
+		parent := current.parentLayer()
+		if parent == nil {
+			break
+		}
+		current = parent
 	}
-	tree.layers = layers
+	tree.base = current.(*diskLayer) // panic if it's not a disk layer
+	tree.lookup = newLookup(head, tree.isDescendant)
 }
 
 // get retrieves a layer belonging to the given state root.
@@ -64,6 +85,43 @@ func (tree *layerTree) get(root common.Hash) layer {
 	return tree.layers[root]
 }
 
+// isDescendant returns whether the specified layer with given root is a
+// descendant of a specific ancestor.
+//
+// This function assumes the read lock has been held.
+func (tree *layerTree) isDescendant(root common.Hash, ancestor common.Hash) bool {
+	subset := tree.descendants[ancestor]
+	if subset == nil {
+		return false
+	}
+	_, ok := subset[root]
+	return ok
+}
+
+// fillAncestors identifies the ancestors of the given layer and populates the
+// descendants set. The ancestors include the diff layers below the supplied
+// layer and also the disk layer.
+//
+// This function assumes the write lock has been held.
+func (tree *layerTree) fillAncestors(layer layer) {
+	hash := layer.rootHash()
+	for {
+		parent := layer.parentLayer()
+		if parent == nil {
+			break
+		}
+		layer = parent
+
+		phash := parent.rootHash()
+		subset := tree.descendants[phash]
+		if subset == nil {
+			subset = make(map[common.Hash]struct{})
+			tree.descendants[phash] = subset
+		}
+		subset[hash] = struct{}{}
+	}
+}
+
 // forEach iterates the stored layers inside and applies the
 // given callback on them.
 func (tree *layerTree) forEach(onLayer func(layer)) {
@@ -101,8 +159,16 @@ func (tree *layerTree) add(root common.Hash, parentRoot common.Hash, block uint6
 	l := parent.update(root, parent.stateID()+1, block, newNodeSet(nodes.Flatten()), states)
 
 	tree.lock.Lock()
+	defer tree.lock.Unlock()
+
+	// Link the given layer into the layer set
 	tree.layers[l.rootHash()] = l
-	tree.lock.Unlock()
+
+	// Link the given layer into its ancestors (up to the current disk layer)
+	tree.fillAncestors(l)
+
+	// Link the given layer into the state mutation history
+	tree.lookup.addLayer(l)
 	return nil
 }
 
@@ -127,8 +193,16 @@ func (tree *layerTree) cap(root common.Hash, layers int) error {
 		if err != nil {
 			return err
 		}
-		// Replace the entire layer tree with the flat base
-		tree.layers = map[common.Hash]layer{base.rootHash(): base}
+		tree.base = base
+
+		// Reset the layer tree with the single new disk layer
+		tree.layers = map[common.Hash]layer{
+			base.rootHash(): base,
+		}
+		// Resets the descendants map, since there's only a single disk layer
+		// with no descendants.
+		tree.descendants = make(map[common.Hash]map[common.Hash]struct{})
+		tree.lookup = newLookup(base, tree.isDescendant)
 		return nil
 	}
 	// Dive until we run out of layers or reach the persistent database
@@ -143,6 +217,11 @@ func (tree *layerTree) cap(root common.Hash, layers int) error {
 	}
 	// We're out of layers, flatten anything below, stopping if it's the disk or if
 	// the memory limit is not yet exceeded.
+	var (
+		err      error
+		replaced layer
+		newBase  *diskLayer
+	)
 	switch parent := diff.parentLayer().(type) {
 	case *diskLayer:
 		return nil
@@ -152,14 +231,33 @@ func (tree *layerTree) cap(root common.Hash, layers int) error {
 		// parent is linked correctly.
 		diff.lock.Lock()
 
-		base, err := parent.persist(false)
+		// Hold the reference of the original layer being replaced
+		replaced = parent
+
+		// Replace the original parent layer with new disk layer. The procedure
+		// can be illustrated as below:
+		//
+		// Before change:
+		//     Chain:
+		//        C1->C2->C3->C4 (HEAD)
+		//          ->C2'->C3'->C4'
+		//
+		// After change:
+		//     Chain:
+		//        (a) C3->C4 (HEAD)
+		//		  (b) C1->C2
+		//		        ->C2'->C3'->C4'
+		// The original C3 is replaced by the new base (with root C3)
+		// Dangling layers in (b) will be removed later
+		newBase, err = parent.persist(false)
 		if err != nil {
 			diff.lock.Unlock()
 			return err
 		}
-		tree.layers[base.rootHash()] = base
-		diff.parent = base
+		tree.layers[newBase.rootHash()] = newBase
 
+		// Link the new parent and release the lock
+		diff.parent = newBase
 		diff.lock.Unlock()
 
 	default:
@@ -173,19 +271,28 @@ func (tree *layerTree) cap(root common.Hash, layers int) error {
 			children[parent] = append(children[parent], root)
 		}
 	}
+	clearDiff := func(layer layer) {
+		diff, ok := layer.(*diffLayer)
+		if !ok {
+			return
+		}
+		tree.lookup.removeLayer(diff)
+	}
 	var remove func(root common.Hash)
 	remove = func(root common.Hash) {
+		clearDiff(tree.layers[root])
+
+		// Unlink the layer from the layer tree and cascade to its children
+		delete(tree.descendants, root)
 		delete(tree.layers, root)
 		for _, child := range children[root] {
 			remove(child)
 		}
 		delete(children, root)
 	}
-	for root, layer := range tree.layers {
-		if dl, ok := layer.(*diskLayer); ok && dl.isStale() {
-			remove(root)
-		}
-	}
+	remove(tree.base.rootHash()) // remove the old/stale disk layer
+	clearDiff(replaced)          // remove the lookup data of the stale parent being replaced
+	tree.base = newBase          // update the base layer with newly constructed one
 	return nil
 }
 
@@ -194,17 +301,41 @@ func (tree *layerTree) bottom() *diskLayer {
 	tree.lock.RLock()
 	defer tree.lock.RUnlock()
 
-	if len(tree.layers) == 0 {
-		return nil // Shouldn't happen, empty tree
+	return tree.base
+}
+
+// lookupAccount returns the layer that is guaranteed to contain the account data
+// corresponding to the specified state root being queried.
+func (tree *layerTree) lookupAccount(accountHash common.Hash, state common.Hash) (layer, error) {
+	// Hold the read lock to prevent the unexpected layer changes
+	tree.lock.RLock()
+	defer tree.lock.RUnlock()
+
+	tip := tree.lookup.accountTip(accountHash, state, tree.base.root)
+	if tip == (common.Hash{}) {
+		return nil, fmt.Errorf("[%#x] %w", state, errSnapshotStale)
 	}
-	// pick a random one as the entry point
-	var current layer
-	for _, layer := range tree.layers {
-		current = layer
-		break
+	l := tree.layers[tip]
+	if l == nil {
+		return nil, fmt.Errorf("triedb layer [%#x] missing", tip)
 	}
-	for current.parentLayer() != nil {
-		current = current.parentLayer()
+	return l, nil
+}
+
+// lookupStorage returns the layer that is guaranteed to contain the storage slot
+// data corresponding to the specified state root being queried.
+func (tree *layerTree) lookupStorage(accountHash common.Hash, slotHash common.Hash, state common.Hash) (layer, error) {
+	// Hold the read lock to prevent the unexpected layer changes
+	tree.lock.RLock()
+	defer tree.lock.RUnlock()
+
+	tip := tree.lookup.storageTip(accountHash, slotHash, state, tree.base.root)
+	if tip == (common.Hash{}) {
+		return nil, fmt.Errorf("[%#x] %w", state, errSnapshotStale)
+	}
+	l := tree.layers[tip]
+	if l == nil {
+		return nil, fmt.Errorf("triedb layer [%#x] missing", tip)
 	}
-	return current.(*diskLayer)
+	return l, nil
 }