pathdb: lookup account and storage nodes

Signed-off-by: jsvisa <delweng@gmail.com>

triedb/pathdb/layertree.go

@@ -338,3 +338,19 @@ func (tree *layerTree) lookupStorage(accountHash common.Hash, slotHash common.Ha
 	}
 	return l, nil
 }
+
+func (tree *layerTree) lookupNode(accountHash common.Hash, path string, 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.nodeTip(accountHash, path, 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 l, nil
+}

triedb/pathdb/lookup.go

@@ -22,6 +22,7 @@ import (
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/log"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -48,6 +49,9 @@ type lookup struct {
 	// where the slot was modified, with the order from oldest to newest.
 	storages map[[64]byte][]common.Hash
 
+	accountNodes map[string][]common.Hash
+	storageNodes map[string][]common.Hash
+
 	// descendant is the callback indicating whether the layer with
 	// given root is a descendant of the one specified by `ancestor`.
 	descendant func(state common.Hash, ancestor common.Hash) bool
@@ -66,6 +70,8 @@ func newLookup(head layer, descendant func(state common.Hash, ancestor common.Ha
 	l := &lookup{
 		accounts:     make(map[common.Hash][]common.Hash),
 		storages:     make(map[[64]byte][]common.Hash),
+		accountNodes: make(map[string][]common.Hash),
+		storageNodes: make(map[string][]common.Hash),
 		descendant:   descendant,
 	}
 	// Apply the diff layers from bottom to top
@@ -161,6 +167,32 @@ func (l *lookup) storageTip(accountHash common.Hash, slotHash common.Hash, state
 	return common.Hash{}
 }
 
+func (l *lookup) nodeTip(accountHash common.Hash, path string, stateID common.Hash, base common.Hash) common.Hash {
+	var list []common.Hash
+	if accountHash == (common.Hash{}) {
+		list = l.accountNodes[path]
+	} else {
+		list = l.storageNodes[accountHash.Hex()+path]
+	}
+	for i := len(list) - 1; i >= 0; i-- {
+		// If the current state matches the stateID, or the requested state is a
+		// descendant of it, return the current state as the most recent one
+		// containing the modified data. Otherwise, the current state may be ahead
+		// of the requested one or belong to a different branch.
+		if list[i] == stateID || l.descendant(stateID, list[i]) {
+			return list[i]
+		}
+	}
+	// No layer matching the stateID or its descendants was found. Use the
+	// current disk layer as a fallback.
+	if base == stateID || l.descendant(stateID, base) {
+		return base
+	}
+	// The layer associated with 'stateID' is not the descendant of the current
+	// disk layer, it's already stale, return nothing.
+	return common.Hash{}
+}
+
 // addLayer traverses the state data retained in the specified diff layer and
 // integrates it into the lookup set.
 //
@@ -176,9 +208,12 @@ func (l *lookup) addLayer(diff *diffLayer) {
 		wg    sync.WaitGroup
 		state = diff.rootHash()
 	)
+	st00 := time.Now()
+	var st0, st1, st2, st3 time.Duration
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
+		st := time.Now()
 		for accountHash := range diff.states.accountData {
 			list, exists := l.accounts[accountHash]
 			if !exists {
@@ -187,11 +222,13 @@ func (l *lookup) addLayer(diff *diffLayer) {
 			list = append(list, state)
 			l.accounts[accountHash] = list
 		}
+		st0 = time.Since(st)
 	}()
 
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
+		st := time.Now()
 		for accountHash, slots := range diff.states.storageData {
 			for slotHash := range slots {
 				key := storageKey(accountHash, slotHash)
@@ -203,8 +240,44 @@ func (l *lookup) addLayer(diff *diffLayer) {
 				l.storages[key] = list
 			}
 		}
+		st1 = time.Since(st)
 	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		st := time.Now()
+		for path := range diff.nodes.accountNodes {
+			list, exists := l.accountNodes[path]
+			if !exists {
+				list = make([]common.Hash, 0, 16) // TODO(rjl493456442) use sync pool
+			}
+			list = append(list, state)
+			l.accountNodes[path] = list
+		}
+		st2 = time.Since(st)
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		st := time.Now()
+		for accountHash, slots := range diff.nodes.storageNodes {
+			for path := range slots {
+				key := accountHash.Hex() + path
+				list, exists := l.storageNodes[key]
+				if !exists {
+					list = make([]common.Hash, 0, 16) // TODO(rjl493456442) use sync pool
+				}
+				list = append(list, state)
+				l.storageNodes[key] = list
+			}
+		}
+		st3 = time.Since(st)
+	}()
+
 	wg.Wait()
+	log.Info("PathDB lookup", "id", diff.id, "block", diff.block, "st0", st0, "st1", st1, "st2", st2, "st3", st3, "elapsed", time.Since(st00))
 }
 
 // removeFromList removes the specified element from the provided list.
@@ -274,5 +347,38 @@ func (l *lookup) removeLayer(diff *diffLayer) error {
 		}
 		return nil
 	})
+
+	eg.Go(func() error {
+		for path := range diff.nodes.accountNodes {
+			found, list := removeFromList(l.accountNodes[path], state)
+			if !found {
+				return fmt.Errorf("account lookup is not found, %x, state: %x", path, state)
+			}
+			if len(list) != 0 {
+				l.accountNodes[path] = list
+			} else {
+				delete(l.accountNodes, path)
+			}
+		}
+		return nil
+	})
+
+	eg.Go(func() error {
+		for accountHash, slots := range diff.nodes.storageNodes {
+			for path := range slots {
+				key := accountHash.Hex() + path
+				found, list := removeFromList(l.storageNodes[key], state)
+				if !found {
+					return fmt.Errorf("storage lookup is not found, %x %x, state: %x", accountHash, path, state)
+				}
+				if len(list) != 0 {
+					l.storageNodes[key] = list
+				} else {
+					delete(l.storageNodes, key)
+				}
+			}
+		}
+		return nil
+	})
 	return eg.Wait()
 }

triedb/pathdb/lookup_test.go

@@ -0,0 +1,123 @@
+package pathdb
+
+import (
+	"crypto/rand"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/trie/trienode"
+)
+
+// generateRandomAccountNodes creates a map of random trie nodes
+func generateRandomAccountNodes(count int) map[string]*trienode.Node {
+	nodes := make(map[string]*trienode.Node, count)
+
+	for i := 0; i < count; i++ {
+		path := make([]byte, 32)
+		rand.Read(path)
+
+		blob := make([]byte, 64)
+		rand.Read(blob)
+
+		var hash common.Hash
+		rand.Read(hash[:])
+
+		nodes[common.Bytes2Hex(path)] = &trienode.Node{Hash: hash, Blob: blob}
+	}
+
+	return nodes
+}
+
+// generateRandomStorageNodes creates a map of storage nodes organized by account
+func generateRandomStorageNodes(accountCount, nodesPerAccount int) map[common.Hash]map[string]*trienode.Node {
+	storageNodes := make(map[common.Hash]map[string]*trienode.Node, accountCount)
+
+	for i := 0; i < accountCount; i++ {
+		var hash common.Hash
+		rand.Read(hash[:])
+
+		storageNodes[hash] = generateRandomAccountNodes(nodesPerAccount)
+	}
+
+	return storageNodes
+}
+
+// addNodes is a helper method for testing that adds nodes to the lookup structure
+func (l *lookup) addNodes(stateHash common.Hash, accountNodes map[string]*trienode.Node, storageNodes map[common.Hash]map[string]*trienode.Node) {
+	// Add account nodes
+	for path := range accountNodes {
+		list, exists := l.accountNodes[path]
+		if !exists {
+			list = make([]common.Hash, 0, 16)
+		}
+		list = append(list, stateHash)
+		l.accountNodes[path] = list
+	}
+
+	// Add storage nodes
+	for accountHash, slots := range storageNodes {
+		for path := range slots {
+			key := accountHash.Hex() + path
+			list, exists := l.storageNodes[key]
+			if !exists {
+				list = make([]common.Hash, 0, 16)
+			}
+			list = append(list, stateHash)
+			l.storageNodes[key] = list
+		}
+	}
+}
+
+func BenchmarkAddNodes(b *testing.B) {
+	tests := []struct {
+		name                string
+		accountNodeCount    int
+		storageAccountCount int
+		nodesPerAccount     int
+	}{
+		{
+			name:                "Small-100-accounts-10-nodes",
+			accountNodeCount:    100,
+			storageAccountCount: 100,
+			nodesPerAccount:     10,
+		},
+		{
+			name:                "Medium-500-accounts-20-nodes",
+			accountNodeCount:    500,
+			storageAccountCount: 500,
+			nodesPerAccount:     20,
+		},
+		{
+			name:                "Large-2000-accounts-40-nodes",
+			accountNodeCount:    2000,
+			storageAccountCount: 2000,
+			nodesPerAccount:     40,
+		},
+	}
+
+	for _, tc := range tests {
+		b.Run(tc.name, func(b *testing.B) {
+			accountNodes := generateRandomAccountNodes(tc.accountNodeCount)
+			storageNodes := generateRandomStorageNodes(tc.storageAccountCount, tc.nodesPerAccount)
+
+			lookup := &lookup{
+				accountNodes: make(map[string][]common.Hash),
+				storageNodes: make(map[string][]common.Hash),
+			}
+
+			var stateHash common.Hash
+			rand.Read(stateHash[:])
+
+			b.ResetTimer()
+			b.ReportAllocs()
+
+			for i := 0; i < b.N; i++ {
+				// Clear the nodes maps for each iteration
+				lookup.accountNodes = make(map[string][]common.Hash)
+				lookup.storageNodes = make(map[string][]common.Hash)
+
+				lookup.addNodes(stateHash, accountNodes, storageNodes)
+			}
+		})
+	}
+}

triedb/pathdb/reader.go

@@ -64,7 +64,14 @@ type reader struct {
 // node info. Don't modify the returned byte slice since it's not deep-copied
 // and still be referenced by database.
 func (r *reader) Node(owner common.Hash, path []byte, hash common.Hash) ([]byte, error) {
-	blob, got, loc, err := r.layer.node(owner, path, 0)
+	l, err := r.db.tree.lookupNode(owner, string(path), r.state)
+	if err != nil {
+		return nil, err
+	}
+	blob, got, loc, err := l.node(owner, path, 0)
+	if errors.Is(err, errSnapshotStale) {
+		blob, got, loc, err = r.layer.node(owner, path, 0)
+	}
 	if err != nil {
 		return nil, err
 	}