all: rework trie and trie committer

core/state/database.go

@@ -88,9 +88,10 @@ type Trie interface {
 	// can be used even if the trie doesn't have one.
 	Hash() common.Hash
 
-	// Commit writes all nodes to the trie's memory database, tracking the internal
-	// and external (for account tries) references.
-	Commit(onleaf trie.LeafCallback) (common.Hash, int, error)
+	// Commit collects all dirty nodes in the trie and replace them with the
+	// corresponding node hash. All collected nodes(including dirty leaves if
+	// collectLeaf is true) will be encapsulated into a nodeset for return.
+	Commit(collectLeaf bool) (common.Hash, *trie.NodeSet, error)
 
 	// NodeIterator returns an iterator that returns nodes of the trie. Iteration
 	// starts at the key after the given start key.

core/state/snapshot/generate.go

@@ -367,7 +367,10 @@ func (dl *diskLayer) generateRange(ctx *generatorContext, owner common.Hash, roo
 		for i, key := range result.keys {
 			snapTrie.Update(key, result.vals[i])
 		}
-		root, _, _ := snapTrie.Commit(nil)
+		root, nodes, _ := snapTrie.Commit(false)
+		if nodes != nil {
+			snapTrieDb.Update(trie.MergeSingle(nodes))
+		}
 		snapTrieDb.Commit(root, false, nil)
 	}
 	// Construct the trie for state iteration, reuse the trie

core/state/snapshot/generate_test.go

@@ -143,6 +143,7 @@ type testHelper struct {
 	diskdb  ethdb.Database
 	triedb  *trie.Database
 	accTrie *trie.SecureTrie
+	nodes   *trie.MergedNodeSet
 }
 
 func newHelper() *testHelper {
@@ -153,6 +154,7 @@ func newHelper() *testHelper {
 		diskdb:  diskdb,
 		triedb:  triedb,
 		accTrie: accTrie,
+		nodes:   trie.NewMergedNodeSet(),
 	}
 }
 
@@ -184,17 +186,22 @@ func (t *testHelper) makeStorageTrie(stateRoot, owner common.Hash, keys []string
 	for i, k := range keys {
 		stTrie.Update([]byte(k), []byte(vals[i]))
 	}
-	var root common.Hash
 	if !commit {
-		root = stTrie.Hash()
-	} else {
-		root, _, _ = stTrie.Commit(nil)
+		return stTrie.Hash().Bytes()
+	}
+	root, nodes, _ := stTrie.Commit(false)
+	if nodes != nil {
+		t.nodes.Merge(nodes)
 	}
 	return root.Bytes()
 }
 
 func (t *testHelper) Commit() common.Hash {
-	root, _, _ := t.accTrie.Commit(nil)
+	root, nodes, _ := t.accTrie.Commit(false)
+	if nodes != nil {
+		t.nodes.Merge(nodes)
+	}
+	t.triedb.Update(t.nodes)
 	t.triedb.Commit(root, false, nil)
 	return root
 }
@@ -378,7 +385,7 @@ func TestGenerateCorruptAccountTrie(t *testing.T) {
 	helper.addTrieAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}) // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
 	helper.addTrieAccount("acc-3", &Account{Balance: big.NewInt(3), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}) // 0x19ead688e907b0fab07176120dceec244a72aff2f0aa51e8b827584e378772f4
 
-	root, _, _ := helper.accTrie.Commit(nil) // Root: 0xa04693ea110a31037fb5ee814308a6f1d76bdab0b11676bdf4541d2de55ba978
+	root, _, _ := helper.accTrie.Commit(false) // Root: 0xa04693ea110a31037fb5ee814308a6f1d76bdab0b11676bdf4541d2de55ba978
 
 	// Delete an account trie leaf and ensure the generator chokes
 	helper.triedb.Commit(root, false, nil)
@@ -413,7 +420,7 @@ func TestGenerateMissingStorageTrie(t *testing.T) {
 	helper.addTrieAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})                                     // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
 	stRoot = helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 	helper.addTrieAccount("acc-3", &Account{Balance: big.NewInt(3), Root: stRoot, CodeHash: emptyCode.Bytes()}) // 0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2
-	root, _, _ := helper.accTrie.Commit(nil)
+	root, _, _ := helper.accTrie.Commit(false)
 
 	// We can only corrupt the disk database, so flush the tries out
 	helper.triedb.Reference(
@@ -458,7 +465,7 @@ func TestGenerateCorruptStorageTrie(t *testing.T) {
 	stRoot = helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 	helper.addTrieAccount("acc-3", &Account{Balance: big.NewInt(3), Root: stRoot, CodeHash: emptyCode.Bytes()}) // 0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2
 
-	root, _, _ := helper.accTrie.Commit(nil)
+	root, _, _ := helper.accTrie.Commit(false)
 
 	// We can only corrupt the disk database, so flush the tries out
 	helper.triedb.Reference(
@@ -825,10 +832,12 @@ func populateDangling(disk ethdb.KeyValueStore) {
 // This test will populate some dangling storages to see if they can be cleaned up.
 func TestGenerateCompleteSnapshotWithDanglingStorage(t *testing.T) {
 	var helper = newHelper()
-	stRoot := helper.makeStorageTrie(common.Hash{}, common.Hash{}, []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 
+	stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-1")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 	helper.addAccount("acc-1", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 	helper.addAccount("acc-2", &Account{Balance: big.NewInt(1), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})
+
+	helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 	helper.addAccount("acc-3", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 
 	helper.addSnapStorage("acc-1", []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"})
@@ -858,10 +867,12 @@ func TestGenerateCompleteSnapshotWithDanglingStorage(t *testing.T) {
 // This test will populate some dangling storages to see if they can be cleaned up.
 func TestGenerateBrokenSnapshotWithDanglingStorage(t *testing.T) {
 	var helper = newHelper()
-	stRoot := helper.makeStorageTrie(common.Hash{}, common.Hash{}, []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 
+	stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-1")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 	helper.addTrieAccount("acc-1", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 	helper.addTrieAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})
+
+	helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 	helper.addTrieAccount("acc-3", &Account{Balance: big.NewInt(3), Root: stRoot, CodeHash: emptyCode.Bytes()})
 
 	populateDangling(helper.diskdb)

core/state/state_object.go

@@ -28,6 +28,7 @@ import (
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/metrics"
 	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/ethereum/go-ethereum/trie"
 )
 
 var emptyCodeHash = crypto.Keccak256(nil)
@@ -375,23 +376,23 @@ func (s *stateObject) updateRoot(db Database) {
 
 // CommitTrie the storage trie of the object to db.
 // This updates the trie root.
-func (s *stateObject) CommitTrie(db Database) (int, error) {
+func (s *stateObject) CommitTrie(db Database) (*trie.NodeSet, error) {
 	// If nothing changed, don't bother with hashing anything
 	if s.updateTrie(db) == nil {
-		return 0, nil
+		return nil, nil
 	}
 	if s.dbErr != nil {
-		return 0, s.dbErr
+		return nil, s.dbErr
 	}
 	// Track the amount of time wasted on committing the storage trie
 	if metrics.EnabledExpensive {
 		defer func(start time.Time) { s.db.StorageCommits += time.Since(start) }(time.Now())
 	}
-	root, committed, err := s.trie.Commit(nil)
+	root, nodes, err := s.trie.Commit(false)
 	if err == nil {
 		s.data.Root = root
 	}
-	return committed, err
+	return nodes, err
 }
 
 // AddBalance adds amount to s's balance.

core/state/statedb.go

@@ -907,7 +907,11 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 	s.IntermediateRoot(deleteEmptyObjects)
 
 	// Commit objects to the trie, measuring the elapsed time
-	var storageCommitted int
+	var (
+		accounts int
+		storages int
+		nodes    = trie.NewMergedNodeSet()
+	)
 	codeWriter := s.db.TrieDB().DiskDB().NewBatch()
 	for addr := range s.stateObjectsDirty {
 		if obj := s.stateObjects[addr]; !obj.deleted {
@@ -917,11 +921,17 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 				obj.dirtyCode = false
 			}
 			// Write any storage changes in the state object to its storage trie
-			committed, err := obj.CommitTrie(s.db)
+			set, err := obj.CommitTrie(s.db)
 			if err != nil {
 				return common.Hash{}, err
 			}
-			storageCommitted += committed
+			// Merge the dirty nodes of storage trie into global set
+			if set != nil {
+				if err := nodes.Merge(set); err != nil {
+					return common.Hash{}, err
+				}
+				storages += set.Len()
+			}
 		}
 	}
 	if len(s.stateObjectsDirty) > 0 {
@@ -937,30 +947,26 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 	if metrics.EnabledExpensive {
 		start = time.Now()
 	}
-	// The onleaf func is called _serially_, so we can reuse the same account
-	// for unmarshalling every time.
-	var account types.StateAccount
-	root, accountCommitted, err := s.trie.Commit(func(_ [][]byte, _ []byte, leaf []byte, parent common.Hash, _ []byte) error {
-		if err := rlp.DecodeBytes(leaf, &account); err != nil {
-			return nil
-		}
-		if account.Root != emptyRoot {
-			s.db.TrieDB().Reference(account.Root, parent)
-		}
-		return nil
-	})
+	root, set, err := s.trie.Commit(true)
 	if err != nil {
 		return common.Hash{}, err
 	}
+	// Merge the dirty nodes of account trie into global set
+	if set != nil {
+		if err := nodes.Merge(set); err != nil {
+			return common.Hash{}, err
+		}
+		accounts = set.Len()
+	}
 	if metrics.EnabledExpensive {
 		s.AccountCommits += time.Since(start)
 
 		accountUpdatedMeter.Mark(int64(s.AccountUpdated))
 		storageUpdatedMeter.Mark(int64(s.StorageUpdated))
 		accountDeletedMeter.Mark(int64(s.AccountDeleted))
 		storageDeletedMeter.Mark(int64(s.StorageDeleted))
-		accountCommittedMeter.Mark(int64(accountCommitted))
-		storageCommittedMeter.Mark(int64(storageCommitted))
+		accountCommittedMeter.Mark(int64(accounts))
+		storageCommittedMeter.Mark(int64(storages))
 		s.AccountUpdated, s.AccountDeleted = 0, 0
 		s.StorageUpdated, s.StorageDeleted = 0, 0
 	}
@@ -984,6 +990,9 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 		}
 		s.snap, s.snapDestructs, s.snapAccounts, s.snapStorage = nil, nil, nil, nil
 	}
+	if err := s.db.TrieDB().Update(nodes); err != nil {
+		return common.Hash{}, err
+	}
 	s.originalRoot = root
 	return root, err
 }

eth/protocols/snap/sync_test.go

@@ -1364,7 +1364,7 @@ func makeAccountTrieNoStorage(n int) (*trie.Trie, entrySlice) {
 		entries = append(entries, elem)
 	}
 	sort.Sort(entries)
-	accTrie.Commit(nil)
+	accTrie.Commit(false)
 	return accTrie, entries
 }
 
@@ -1420,7 +1420,7 @@ func makeBoundaryAccountTrie(n int) (*trie.Trie, entrySlice) {
 		entries = append(entries, elem)
 	}
 	sort.Sort(entries)
-	trie.Commit(nil)
+	trie.Commit(false)
 	return trie, entries
 }
 
@@ -1444,7 +1444,7 @@ func makeAccountTrieWithStorageWithUniqueStorage(accounts, slots int, code bool)
 		// Create a storage trie
 		stTrie, stEntries := makeStorageTrieWithSeed(common.BytesToHash(key), uint64(slots), i, db)
 		stRoot := stTrie.Hash()
-		stTrie.Commit(nil)
+		stTrie.Commit(false)
 		value, _ := rlp.EncodeToBytes(&types.StateAccount{
 			Nonce:    i,
 			Balance:  big.NewInt(int64(i)),
@@ -1460,7 +1460,7 @@ func makeAccountTrieWithStorageWithUniqueStorage(accounts, slots int, code bool)
 	}
 	sort.Sort(entries)
 
-	accTrie.Commit(nil)
+	accTrie.Commit(false)
 	return accTrie, entries, storageTries, storageEntries
 }
 
@@ -1491,7 +1491,7 @@ func makeAccountTrieWithStorage(accounts, slots int, code, boundary bool) (*trie
 			stTrie, stEntries = makeStorageTrieWithSeed(common.BytesToHash(key), uint64(slots), 0, db)
 		}
 		stRoot := stTrie.Hash()
-		stTrie.Commit(nil)
+		stTrie.Commit(false)
 
 		value, _ := rlp.EncodeToBytes(&types.StateAccount{
 			Nonce:    i,
@@ -1507,7 +1507,7 @@ func makeAccountTrieWithStorage(accounts, slots int, code, boundary bool) (*trie
 		storageEntries[common.BytesToHash(key)] = stEntries
 	}
 	sort.Sort(entries)
-	accTrie.Commit(nil)
+	accTrie.Commit(false)
 	return accTrie, entries, storageTries, storageEntries
 }
 
@@ -1530,7 +1530,7 @@ func makeStorageTrieWithSeed(owner common.Hash, n, seed uint64, db *trie.Databas
 		entries = append(entries, elem)
 	}
 	sort.Sort(entries)
-	trie.Commit(nil)
+	trie.Commit(false)
 	return trie, entries
 }
 
@@ -1581,7 +1581,7 @@ func makeBoundaryStorageTrie(owner common.Hash, n int, db *trie.Database) (*trie
 		entries = append(entries, elem)
 	}
 	sort.Sort(entries)
-	trie.Commit(nil)
+	trie.Commit(false)
 	return trie, entries
 }
 

light/postprocess.go

@@ -217,10 +217,15 @@ func (c *ChtIndexerBackend) Process(ctx context.Context, header *types.Header) e
 
 // Commit implements core.ChainIndexerBackend
 func (c *ChtIndexerBackend) Commit() error {
-	root, _, err := c.trie.Commit(nil)
+	root, nodes, err := c.trie.Commit(false)
 	if err != nil {
 		return err
 	}
+	if nodes != nil {
+		if err := c.triedb.Update(trie.MergeSingle(nodes)); err != nil {
+			return err
+		}
+	}
 	// Pruning historical trie nodes if necessary.
 	if !c.disablePruning {
 		// Flush the triedb and track the latest trie nodes.
@@ -453,10 +458,15 @@ func (b *BloomTrieIndexerBackend) Commit() error {
 			b.trie.Delete(encKey[:])
 		}
 	}
-	root, _, err := b.trie.Commit(nil)
+	root, nodes, err := b.trie.Commit(false)
 	if err != nil {
 		return err
 	}
+	if nodes != nil {
+		if err := b.triedb.Update(trie.MergeSingle(nodes)); err != nil {
+			return err
+		}
+	}
 	// Pruning historical trie nodes if necessary.
 	if !b.disablePruning {
 		// Flush the triedb and track the latest trie nodes.
@@ -504,4 +514,4 @@ func (b *BloomTrieIndexerBackend) Prune(threshold uint64) error {
 	}
 	log.Debug("Prune history bloombits", "threshold", threshold, "elapsed", common.PrettyDuration(time.Since(start)))
 	return nil
-}
+}

light/trie.go

@@ -137,11 +137,11 @@ func (t *odrTrie) TryDelete(key []byte) error {
 	})
 }
 
-func (t *odrTrie) Commit(onleaf trie.LeafCallback) (common.Hash, int, error) {
+func (t *odrTrie) Commit(collectLeaf bool) (common.Hash, *trie.NodeSet, error) {
 	if t.trie == nil {
-		return t.id.Root, 0, nil
+		return t.id.Root, nil, nil
 	}
-	return t.trie.Commit(onleaf)
+	return t.trie.Commit(collectLeaf)
 }
 
 func (t *odrTrie) Hash() common.Hash {

tests/fuzzers/stacktrie/trie_fuzzer.go

@@ -174,10 +174,13 @@ func (f *fuzzer) fuzz() int {
 		return 0
 	}
 	// Flush trie -> database
-	rootA, _, err := trieA.Commit(nil)
+	rootA, nodes, err := trieA.Commit(false)
 	if err != nil {
 		panic(err)
 	}
+	if nodes != nil {
+		dbA.Update(trie.MergeSingle(nodes))
+	}
 	// Flush memdb -> disk (sponge)
 	dbA.Commit(rootA, false, nil)