triedb/pathdb: introduce lookup structure to optimize state access

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
 }