core/state: implement fast storage deletion (ethereum#27955)

This changes implements faster post-selfdestruct iteration of storage slots for deletion, by using snapshot-storage+stacktrie to recover the trienodes to be deleted. This mechanism is only implemented for path-based schema. 

For hash-based schema, the entire post-selfdestruct storage iteration is skipped, with this change, since hash-based does not actually perform deletion anyway. 

---------

Co-authored-by: Martin Holst Swende <martin@swende.se>

core/state/statedb.go

@@ -36,6 +36,12 @@ import (
 	"github.com/ethereum/go-ethereum/trie/triestate"
 )
 
+const (
+	// storageDeleteLimit denotes the highest permissible memory allocation
+	// employed for contract storage deletion.
+	storageDeleteLimit = 512 * 1024 * 1024
+)
+
 type revision struct {
 	id           int
 	journalIndex int
@@ -983,59 +989,130 @@ func (s *StateDB) clearJournalAndRefund() {
 	s.validRevisions = s.validRevisions[:0] // Snapshots can be created without journal entries
 }
 
-// deleteStorage iterates the storage trie belongs to the account and mark all
-// slots inside as deleted.
-func (s *StateDB) deleteStorage(addr common.Address, addrHash common.Hash, root common.Hash) (bool, map[common.Hash][]byte, *trienode.NodeSet, error) {
-	start := time.Now()
+// fastDeleteStorage is the function that efficiently deletes the storage trie
+// of a specific account. It leverages the associated state snapshot for fast
+// storage iteration and constructs trie node deletion markers by creating
+// stack trie with iterated slots.
+func (s *StateDB) fastDeleteStorage(addrHash common.Hash, root common.Hash) (bool, common.StorageSize, map[common.Hash][]byte, *trienode.NodeSet, error) {
+	iter, err := s.snaps.StorageIterator(s.originalRoot, addrHash, common.Hash{})
+	if err != nil {
+		return false, 0, nil, nil, err
+	}
+	defer iter.Release()
+
+	var (
+		size  common.StorageSize
+		nodes = trienode.NewNodeSet(addrHash)
+		slots = make(map[common.Hash][]byte)
+	)
+	stack := trie.NewStackTrie(func(owner common.Hash, path []byte, hash common.Hash, blob []byte) {
+		nodes.AddNode(path, trienode.NewDeleted())
+		size += common.StorageSize(len(path))
+	})
+	for iter.Next() {
+		if size > storageDeleteLimit {
+			return true, size, nil, nil, nil
+		}
+		slot := common.CopyBytes(iter.Slot())
+		if iter.Error() != nil { // error might occur after Slot function
+			return false, 0, nil, nil, err
+		}
+		size += common.StorageSize(common.HashLength + len(slot))
+		slots[iter.Hash()] = slot
+
+		if err := stack.Update(iter.Hash().Bytes(), slot); err != nil {
+			return false, 0, nil, nil, err
+		}
+	}
+	if iter.Error() != nil { // error might occur during iteration
+		return false, 0, nil, nil, err
+	}
+	if stack.Hash() != root {
+		return false, 0, nil, nil, fmt.Errorf("snapshot is not matched, exp %x, got %x", root, stack.Hash())
+	}
+	return false, size, slots, nodes, nil
+}
+
+// slowDeleteStorage serves as a less-efficient alternative to "fastDeleteStorage,"
+// employed when the associated state snapshot is not available. It iterates the
+// storage slots along with all internal trie nodes via trie directly.
+func (s *StateDB) slowDeleteStorage(addr common.Address, addrHash common.Hash, root common.Hash) (bool, common.StorageSize, map[common.Hash][]byte, *trienode.NodeSet, error) {
 	tr, err := s.db.OpenStorageTrie(s.originalRoot, addr, root)
 	if err != nil {
-		return false, nil, nil, fmt.Errorf("failed to open storage trie, err: %w", err)
+		return false, 0, nil, nil, fmt.Errorf("failed to open storage trie, err: %w", err)
 	}
 	it, err := tr.NodeIterator(nil)
 	if err != nil {
-		return false, nil, nil, fmt.Errorf("failed to open storage iterator, err: %w", err)
+		return false, 0, nil, nil, fmt.Errorf("failed to open storage iterator, err: %w", err)
 	}
 	var (
-		set       = trienode.NewNodeSet(addrHash)
-		slots     = make(map[common.Hash][]byte)
-		stateSize common.StorageSize
-		nodeSize  common.StorageSize
+		size  common.StorageSize
+		nodes = trienode.NewNodeSet(addrHash)
+		slots = make(map[common.Hash][]byte)
 	)
 	for it.Next(true) {
-		// arbitrary stateSize limit, make it configurable
-		if stateSize+nodeSize > 512*1024*1024 {
-			log.Info("Skip large storage deletion", "address", addr.Hex(), "states", stateSize, "nodes", nodeSize)
-			if metrics.EnabledExpensive {
-				slotDeletionSkip.Inc(1)
-			}
-			return true, nil, nil, nil
+		if size > storageDeleteLimit {
+			return true, size, nil, nil, nil
 		}
 		if it.Leaf() {
 			slots[common.BytesToHash(it.LeafKey())] = common.CopyBytes(it.LeafBlob())
-			stateSize += common.StorageSize(common.HashLength + len(it.LeafBlob()))
+			size += common.StorageSize(common.HashLength + len(it.LeafBlob()))
 			continue
 		}
 		if it.Hash() == (common.Hash{}) {
 			continue
 		}
-		nodeSize += common.StorageSize(len(it.Path()))
-		set.AddNode(it.Path(), trienode.NewDeleted())
+		size += common.StorageSize(len(it.Path()))
+		nodes.AddNode(it.Path(), trienode.NewDeleted())
 	}
 	if err := it.Error(); err != nil {
+		return false, 0, nil, nil, err
+	}
+	return false, size, slots, nodes, nil
+}
+
+// deleteStorage is designed to delete the storage trie of a designated account.
+// It could potentially be terminated if the storage size is excessively large,
+// potentially leading to an out-of-memory panic. The function will make an attempt
+// to utilize an efficient strategy if the associated state snapshot is reachable;
+// otherwise, it will resort to a less-efficient approach.
+func (s *StateDB) deleteStorage(addr common.Address, addrHash common.Hash, root common.Hash) (bool, map[common.Hash][]byte, *trienode.NodeSet, error) {
+	var (
+		start   = time.Now()
+		err     error
+		aborted bool
+		size    common.StorageSize
+		slots   map[common.Hash][]byte
+		nodes   *trienode.NodeSet
+	)
+	// The fast approach can be failed if the snapshot is not fully
+	// generated, or it's internally corrupted. Fallback to the slow
+	// one just in case.
+	if s.snap != nil {
+		aborted, size, slots, nodes, err = s.fastDeleteStorage(addrHash, root)
+	}
+	if s.snap == nil || err != nil {
+		aborted, size, slots, nodes, err = s.slowDeleteStorage(addr, addrHash, root)
+	}
+	if err != nil {
 		return false, nil, nil, err
 	}
 	if metrics.EnabledExpensive {
-		if int64(len(slots)) > slotDeletionMaxCount.Value() {
-			slotDeletionMaxCount.Update(int64(len(slots)))
+		if aborted {
+			slotDeletionSkip.Inc(1)
 		}
-		if int64(stateSize+nodeSize) > slotDeletionMaxSize.Value() {
-			slotDeletionMaxSize.Update(int64(stateSize + nodeSize))
+		n := int64(len(slots))
+		if n > slotDeletionMaxCount.Value() {
+			slotDeletionMaxCount.Update(n)
+		}
+		if int64(size) > slotDeletionMaxSize.Value() {
+			slotDeletionMaxSize.Update(int64(size))
 		}
 		slotDeletionTimer.UpdateSince(start)
-		slotDeletionCount.Mark(int64(len(slots)))
-		slotDeletionSize.Mark(int64(stateSize + nodeSize))
+		slotDeletionCount.Mark(n)
+		slotDeletionSize.Mark(int64(size))
 	}
-	return false, slots, set, nil
+	return aborted, slots, nodes, nil
 }
 
 // handleDestruction processes all destruction markers and deletes the account
@@ -1063,7 +1140,13 @@ func (s *StateDB) deleteStorage(addr common.Address, addrHash common.Hash, root
 // In case (d), **original** account along with its storages should be deleted,
 // with their values be tracked as original value.
 func (s *StateDB) handleDestruction(nodes *trienode.MergedNodeSet) (map[common.Address]struct{}, error) {
+	// Short circuit if geth is running with hash mode. This procedure can consume
+	// considerable time and storage deletion isn't supported in hash mode, thus
+	// preemptively avoiding unnecessary expenses.
 	incomplete := make(map[common.Address]struct{})
+	if s.db.TrieDB().Scheme() == rawdb.HashScheme {
+		return incomplete, nil
+	}
 	for addr, prev := range s.stateObjectsDestruct {
 		// The original account was non-existing, and it's marked as destructed
 		// in the scope of block. It can be case (a) or (b).

core/state/statedb_fuzz_test.go

@@ -31,10 +31,12 @@ import (
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/core/state/snapshot"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
+	"github.com/ethereum/go-ethereum/trie/triedb/pathdb"
 	"github.com/ethereum/go-ethereum/trie/triestate"
 )
 
@@ -179,16 +181,28 @@ func (test *stateTest) run() bool {
 			storageList = append(storageList, copy2DSet(states.Storages))
 		}
 		disk      = rawdb.NewMemoryDatabase()
-		tdb       = trie.NewDatabase(disk, &trie.Config{OnCommit: onCommit})
+		tdb       = trie.NewDatabase(disk, &trie.Config{OnCommit: onCommit, PathDB: pathdb.Defaults})
 		sdb       = NewDatabaseWithNodeDB(disk, tdb)
 		byzantium = rand.Intn(2) == 0
 	)
+	defer disk.Close()
+	defer tdb.Close()
+
+	var snaps *snapshot.Tree
+	if rand.Intn(3) == 0 {
+		snaps, _ = snapshot.New(snapshot.Config{
+			CacheSize:  1,
+			Recovery:   false,
+			NoBuild:    false,
+			AsyncBuild: false,
+		}, disk, tdb, types.EmptyRootHash)
+	}
 	for i, actions := range test.actions {
 		root := types.EmptyRootHash
 		if i != 0 {
 			root = roots[len(roots)-1]
 		}
-		state, err := New(root, sdb, nil)
+		state, err := New(root, sdb, snaps)
 		if err != nil {
 			panic(err)
 		}

core/state/statedb_test.go

@@ -39,6 +39,8 @@ import (
 	"github.com/ethereum/go-ethereum/trie"
 	"github.com/ethereum/go-ethereum/trie/triedb/hashdb"
 	"github.com/ethereum/go-ethereum/trie/triedb/pathdb"
+	"github.com/ethereum/go-ethereum/trie/trienode"
+	"github.com/holiman/uint256"
 )
 
 // Tests that updating a state trie does not leak any database writes prior to
@@ -1135,3 +1137,57 @@ func TestResetObject(t *testing.T) {
 		t.Fatalf("Unexpected storage slot value %v", slot)
 	}
 }
+
+func TestDeleteStorage(t *testing.T) {
+	var (
+		disk     = rawdb.NewMemoryDatabase()
+		tdb      = trie.NewDatabase(disk, nil)
+		db       = NewDatabaseWithNodeDB(disk, tdb)
+		snaps, _ = snapshot.New(snapshot.Config{CacheSize: 10}, disk, tdb, types.EmptyRootHash)
+		state, _ = New(types.EmptyRootHash, db, snaps)
+		addr     = common.HexToAddress("0x1")
+	)
+	// Initialize account and populate storage
+	state.SetBalance(addr, big.NewInt(1))
+	state.CreateAccount(addr)
+	for i := 0; i < 1000; i++ {
+		slot := common.Hash(uint256.NewInt(uint64(i)).Bytes32())
+		value := common.Hash(uint256.NewInt(uint64(10 * i)).Bytes32())
+		state.SetState(addr, slot, value)
+	}
+	root, _ := state.Commit(0, true)
+	// Init phase done, create two states, one with snap and one without
+	fastState, _ := New(root, db, snaps)
+	slowState, _ := New(root, db, nil)
+
+	obj := fastState.GetOrNewStateObject(addr)
+	storageRoot := obj.data.Root
+
+	_, _, fastNodes, err := fastState.deleteStorage(addr, crypto.Keccak256Hash(addr[:]), storageRoot)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, _, slowNodes, err := slowState.deleteStorage(addr, crypto.Keccak256Hash(addr[:]), storageRoot)
+	if err != nil {
+		t.Fatal(err)
+	}
+	check := func(set *trienode.NodeSet) string {
+		var a []string
+		set.ForEachWithOrder(func(path string, n *trienode.Node) {
+			if n.Hash != (common.Hash{}) {
+				t.Fatal("delete should have empty hashes")
+			}
+			if len(n.Blob) != 0 {
+				t.Fatal("delete should have have empty blobs")
+			}
+			a = append(a, fmt.Sprintf("%x", path))
+		})
+		return strings.Join(a, ",")
+	}
+	slowRes := check(slowNodes)
+	fastRes := check(fastNodes)
+	if slowRes != fastRes {
+		t.Fatalf("difference found:\nfast: %v\nslow: %v\n", fastRes, slowRes)
+	}
+}