diff --git a/core/state/iterator.go b/core/state/iterator.go index a051e5e01a..84259ac512 100644 --- a/core/state/iterator.go +++ b/core/state/iterator.go @@ -26,6 +26,11 @@ import ( "github.com/ethereum/go-ethereum/trie" ) +// NodeIteratorSubtreeCallback is a callback that is invoked by the trie node +// iterator whenever it descends into a new subtree, giving the possibility of +// skipping the branch if it would be non-useful (e.g. already processed before). +type NodeIteratorSubtreeCallback func(hash, parent common.Hash) bool + // NodeIterator is an iterator to traverse the entire state trie post-order, // including all of the contract code and contract state tries. type NodeIterator struct { @@ -38,6 +43,14 @@ type NodeIterator struct { codeHash common.Hash // Hash of the contract source code code []byte // Source code associated with a contract + // Callback method enabling the user to control subtree descent. + // + // Note: This hook is invoked pre-order to allow the user to cancel the traversal + // of a subtree, but the iterator itself is post-order. This means that the hook + // may be invoked multiple times during a single iteration step (when descending) + // or never in multiple iteration steps (when ascending). + PreOrderHook NodeIteratorSubtreeCallback + Hash common.Hash // Hash of the current entry being iterated (nil if not standalone) Entry interface{} // Current state entry being iterated (internal representation) Parent common.Hash // Hash of the first full ancestor node (nil if current is the root) @@ -66,6 +79,7 @@ func (it *NodeIterator) step() { // Initialize the iterator if we've just started if it.stateIt == nil { it.stateIt = trie.NewNodeIterator(it.state.trie.Trie) + it.stateIt.PreOrderHook = trie.NodeIteratorSubtreeCallback(it.PreOrderHook) } // If we had data nodes previously, we surely have at least state nodes if it.dataIt != nil { @@ -99,22 +113,25 @@ func (it *NodeIterator) step() { if err != nil { panic(err) } + it.accountHash = it.stateIt.Parent + + if bytes.Compare(account.CodeHash, emptyCodeHash) != 0 { + if hash := common.BytesToHash(account.CodeHash); it.PreOrderHook == nil || it.PreOrderHook(hash, it.accountHash) { + it.codeHash = hash + if it.code, err = it.state.db.Get(account.CodeHash); err != nil { + panic(fmt.Sprintf("code %x: %v", account.CodeHash, err)) + } + } + } dataTrie, err := trie.New(account.Root, it.state.db) if err != nil { panic(err) } it.dataIt = trie.NewNodeIterator(dataTrie) + it.dataIt.PreOrderHook = trie.NodeIteratorSubtreeCallback(it.PreOrderHook) if !it.dataIt.Next() { it.dataIt = nil } - if bytes.Compare(account.CodeHash, emptyCodeHash) != 0 { - it.codeHash = common.BytesToHash(account.CodeHash) - it.code, err = it.state.db.Get(account.CodeHash) - if err != nil { - panic(fmt.Sprintf("code %x: %v", account.CodeHash, err)) - } - } - it.accountHash = it.stateIt.Parent } // retrieve pulls and caches the current state entry the iterator is traversing. diff --git a/core/state/iterator_test.go b/core/state/iterator_test.go index d44e15be7d..0e7aba78c6 100644 --- a/core/state/iterator_test.go +++ b/core/state/iterator_test.go @@ -59,3 +59,63 @@ func TestNodeIteratorCoverage(t *testing.T) { } } } + +// Tests that the node iterator hook is invoked for all the nodes of the state, +// also testing that the iterator indeed avoids visiting aborted branches. +func TestNodeIteratorHookCoverageFull(t *testing.T) { testNodeIteratorHookCoverage(t, false) } +func TestNodeIteratorHookCoverageDedup(t *testing.T) { testNodeIteratorHookCoverage(t, true) } + +func testNodeIteratorHookCoverage(t *testing.T, dedup bool) { + // Create some arbitrary test state to iterate + db, root, _ := makeTestState(nil) + + state, err := New(root, db) + if err != nil { + t.Fatalf("failed to create state trie at %x: %v", root, err) + } + // Gather all the node hashes found by the iterator + hashes := make(map[common.Hash]int) + + it := NewNodeIterator(state) + it.PreOrderHook = func(hash, parent common.Hash) bool { + if !dedup { + return true + } + return hashes[hash] == 0 + } + for it.Next() { + if it.Hash != (common.Hash{}) { + hashes[it.Hash]++ + } + } + // Cross check the hashes and the database itself + for hash, _ := range hashes { + if _, err := db.Get(hash.Bytes()); err != nil { + t.Errorf("failed to retrieve reported node %x: %v", hash, err) + } + } + for _, key := range db.(*ethdb.MemDatabase).Keys() { + if bytes.HasPrefix(key, []byte("secure-key-")) { + continue + } + if bytes.HasPrefix(key, trie.ParentReferenceIndexPrefix) { + continue + } + if _, ok := hashes[common.BytesToHash(key)]; !ok { + t.Errorf("state entry not reported %x", key) + } + } + // Check whether duplicates were avoided or not + duplicates := 0 + for hash, count := range hashes { + if count > 1 { + duplicates++ + if dedup { + t.Errorf("duplicate (%d) iteration: %x", count, hash) + } + } + } + if !dedup && duplicates == 0 { + t.Errorf("iterator didn't traverse common subtrees") + } +} diff --git a/core/state/sync_test.go b/core/state/sync_test.go index 30cd70eec6..e975a317ec 100644 --- a/core/state/sync_test.go +++ b/core/state/sync_test.go @@ -44,22 +44,24 @@ func makeTestState(referrers []common.Hash) (ethdb.Database, common.Hash, []*tes // Fill it with some arbitrary data accounts := []*testAccount{} - for i := byte(0); i < 96; i++ { - obj := state.GetOrNewStateObject(common.BytesToAddress([]byte{i})) - acc := &testAccount{address: common.BytesToAddress([]byte{i})} + for i := byte(0); i < 1; i++ { + for j := byte(0); j < 2; j++ { + obj := state.GetOrNewStateObject(common.BytesToAddress([]byte{i, j})) + acc := &testAccount{address: common.BytesToAddress([]byte{i, j})} - obj.AddBalance(big.NewInt(int64(11 * i))) - acc.balance = big.NewInt(int64(11 * i)) + obj.AddBalance(big.NewInt(int64(11 * i))) + acc.balance = big.NewInt(int64(11 * i)) - obj.SetNonce(uint64(42 * i)) - acc.nonce = uint64(42 * i) + obj.SetNonce(uint64(42 * i)) + acc.nonce = uint64(42 * i) - if i%3 == 0 { - obj.SetCode([]byte{i, i, i, i, i}) - acc.code = []byte{i, i, i, i, i} + if i%3 == 0 { + obj.SetCode([]byte{i, i, i, i, i}) + acc.code = []byte{i, i, i, i, i} + } + state.UpdateStateObject(obj) + accounts = append(accounts, acc) } - state.UpdateStateObject(obj) - accounts = append(accounts, acc) } root, _ := state.CommitIndexed(referrers) diff --git a/trie/iterator.go b/trie/iterator.go index 59ac3c5182..df9efece7f 100644 --- a/trie/iterator.go +++ b/trie/iterator.go @@ -159,11 +159,24 @@ type nodeIteratorState struct { child int // Child to be processed next } +// NodeIteratorSubtreeCallback is a callback that is invoked by the trie node +// iterator whenever it descends into a new subtree, giving the possibility of +// skipping the branch if it would be non-useful (e.g. already processed before). +type NodeIteratorSubtreeCallback func(hash, parent common.Hash) bool + // NodeIterator is an iterator to traverse the trie post-order. type NodeIterator struct { trie *Trie // Trie being iterated stack []*nodeIteratorState // Hierarchy of trie nodes persisting the iteration state + // Callback method enabling the user to control subtree descent. + // + // Note: This hook is invoked pre-order to allow the user to cancel the traversal + // of a subtree, but the iterator itself is post-order. This means that the hook + // may be invoked multiple times during a single iteration step (when descending) + // or never in multiple iteration steps (when ascending). + PreOrderHook NodeIteratorSubtreeCallback + Hash common.Hash // Hash of the current node being iterated (nil if not standalone) Node node // Current node being iterated (internal representation) Parent common.Hash // Hash of the first full ancestor node (nil if current is the root) @@ -241,7 +254,9 @@ func (it *NodeIterator) step() { if err != nil { panic(err) } - it.stack = append(it.stack, &nodeIteratorState{hash: common.BytesToHash(hash), node: node, parent: ancestor, child: -1}) + if hash := common.BytesToHash(hash); it.PreOrderHook == nil || it.PreOrderHook(hash, ancestor) { + it.stack = append(it.stack, &nodeIteratorState{hash: hash, node: node, parent: ancestor, child: -1}) + } } else { break } diff --git a/trie/iterator_test.go b/trie/iterator_test.go index b095de855e..dadee3fdc3 100644 --- a/trie/iterator_test.go +++ b/trie/iterator_test.go @@ -79,3 +79,55 @@ func TestNodeIteratorCoverage(t *testing.T) { } } } + +// Tests that the node iterator hook is invoked for all the nodes of the trie, +// also testing that the iterator indeed avoids visiting aborted branches. +func TestNodeIteratorHookCoverageFull(t *testing.T) { testNodeIteratorHookCoverage(t, false) } +func TestNodeIteratorHookCoverageDedup(t *testing.T) { testNodeIteratorHookCoverage(t, true) } + +func testNodeIteratorHookCoverage(t *testing.T, dedup bool) { + // Create some arbitrary test trie to iterate + db, trie, _ := makeTestTrie(nil) + + // Gather all the node hashes found by the iterator + hashes := make(map[common.Hash]int) + + it := NewNodeIterator(trie) + it.PreOrderHook = func(hash, parent common.Hash) bool { + if !dedup { + return true + } + return hashes[hash] == 0 + } + for it.Next() { + if it.Hash != (common.Hash{}) { + hashes[it.Hash]++ + } + } + // Cross check the hashes and the database itself + for hash, _ := range hashes { + if _, err := db.Get(hash.Bytes()); err != nil { + t.Errorf("failed to retrieve reported node %x: %v", hash, err) + } + } + for _, key := range db.(*ethdb.MemDatabase).Keys() { + if len(key) == common.HashLength { + if _, ok := hashes[common.BytesToHash(key)]; !ok { + t.Errorf("state entry not reported %x", key) + } + } + } + // Check whether duplicates were avoided or not + duplicates := 0 + for hash, count := range hashes { + if count > 1 { + duplicates++ + if dedup { + t.Errorf("duplicate (%d) iteration: %x", count, hash) + } + } + } + if !dedup && duplicates == 0 { + t.Errorf("iterator didn't traverse common subtrees") + } +} diff --git a/trie/sync.go b/trie/sync.go index e9955cfb74..0daf0ecb92 100644 --- a/trie/sync.go +++ b/trie/sync.go @@ -34,11 +34,19 @@ type request struct { depth int // Depth level within the trie the node is located to prioritise DFS deps int // Number of dependencies before allowed to commit this node - referrers []common.Hash // External parents of this node to index in addition to internal ones + referrers map[common.Hash]struct{} // External parents of this node to index in addition to internal ones callback TrieSyncLeafCallback // Callback to invoke if a leaf node it reached on this branch } +// referrer adds a new referring parent to a sync node request. +func (r *request) referrer(hash common.Hash) { + if r.referrers == nil { + r.referrers = make(map[common.Hash]struct{}) + } + r.referrers[hash] = struct{}{} +} + // SyncResult is a simple list to return missing nodes along with their request // hashes. type SyncResult struct { @@ -72,14 +80,14 @@ func NewTrieSync(root common.Hash, database ethdb.Database, parent common.Hash, } // AddSubTrie registers a new trie to the sync code, rooted at the designated parent. -func (s *TrieSync) AddSubTrie(root common.Hash, depth int, parent common.Hash, callback TrieSyncLeafCallback) { +func (s *TrieSync) AddSubTrie(root common.Hash, depth int, parent common.Hash, callback TrieSyncLeafCallback) error { // Short circuit if the trie is empty or already known if root == emptyRoot { - return + return nil } blob, _ := s.database.Get(root.Bytes()) if local, err := decodeNode(blob); local != nil && err == nil { - return + return storeParentReferenceEntry(parent.Bytes(), root.Bytes(), s.database) } // Assemble the new sub-trie sync request node := node(hashNode(root.Bytes())) @@ -99,22 +107,23 @@ func (s *TrieSync) AddSubTrie(root common.Hash, depth int, parent common.Hash, c ancestor.deps++ req.parents = append(req.parents, ancestor) } - req.referrers = append(req.referrers, parent) + req.referrer(parent) } s.schedule(req) + return nil } // AddRawEntry schedules the direct retrieval of a state entry that should not be // interpreted as a trie node, but rather accepted and stored into the database // as is. This method's goal is to support misc state metadata retrievals (e.g. // contract code). -func (s *TrieSync) AddRawEntry(hash common.Hash, depth int, parent common.Hash) { +func (s *TrieSync) AddRawEntry(hash common.Hash, depth int, parent common.Hash) error { // Short circuit if the entry is empty or already known if hash == emptyState { - return + return nil } if blob, _ := s.database.Get(hash.Bytes()); blob != nil { - return + return storeParentReferenceEntry(parent.Bytes(), hash.Bytes(), s.database) } // Assemble the new sub-trie sync request req := &request{ @@ -131,9 +140,10 @@ func (s *TrieSync) AddRawEntry(hash common.Hash, depth int, parent common.Hash) ancestor.deps++ req.parents = append(req.parents, ancestor) } - req.referrers = append(req.referrers, parent) + req.referrer(parent) } s.schedule(req) + return nil } // Missing retrieves the known missing nodes from the trie for retrieval. @@ -196,6 +206,13 @@ func (s *TrieSync) schedule(req *request) { // If we're already requesting this node, add a new reference and stop if old, ok := s.requests[req.hash]; ok { old.parents = append(old.parents, req.parents...) + + for hash, _ := range req.referrers { + old.referrer(hash) + } + for _, parent := range req.parents { + old.referrer(parent.hash) + } return } // Schedule the request for future retrieval @@ -280,7 +297,7 @@ func (s *TrieSync) commit(req *request, batch ethdb.Batch) (err error) { return err } } - for _, referrer := range req.referrers { + for referrer, _ := range req.referrers { if err := storeParentReferenceEntry(referrer[:], req.hash[:], batch); err != nil { return err }