Merge d477a72666 into 406a852ec8

cmd/geth/snapshot.go

@@ -18,13 +18,17 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"os"
 	"slices"
+	"sync/atomic"
 	"time"
 
+	"golang.org/x/sync/errgroup"
+
 	"github.com/ethereum/go-ethereum/cmd/utils"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
@@ -33,9 +37,13 @@ import (
 	"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/ethdb"
 	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/node"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
+	"github.com/ethereum/go-ethereum/triedb"
+	"github.com/ethereum/go-ethereum/triedb/database"
 	"github.com/urfave/cli/v2"
 )
 
@@ -105,12 +113,22 @@ information about the specified address.
 				Usage:     "Traverse the state with given root hash and perform quick verification",
 				ArgsUsage: "<root>",
 				Action:    traverseState,
-				Flags:     slices.Concat(utils.NetworkFlags, utils.DatabaseFlags),
+				Flags:     slices.Concat(utils.TraverseStateFlags, utils.NetworkFlags, utils.DatabaseFlags),
 				Description: `
-geth snapshot traverse-state <state-root>
-will traverse the whole state from the given state root and will abort if any
-referenced trie node or contract code is missing. This command can be used for
-state integrity verification. The default checking target is the HEAD state.
+geth snapshot traverse-state [--account <account>] [--start <key>] [--limit <key>] <state-root>
+
+1. Traverse the whole state from the given state root:
+- --start: starting account key (64/66 chars hex) [optional]
+- --limit: ending account key (64/66 chars hex) [optional]
+
+2. Traverse a specific account's storage:
+- --account: account address (40/42 chars) or hash (64/66 chars) [required]
+- --start: starting storage key (64/66 chars hex) [optional]
+- --limit: ending storage key (64/66 chars hex) [optional]
+
+The default checking state root is the HEAD state if not specified.
+The command will abort if any referenced trie node or contract code is missing.
+This can be used for state integrity verification. The default target is HEAD state.
 
 It's also usable without snapshot enabled.
 `,
@@ -120,13 +138,24 @@ It's also usable without snapshot enabled.
 				Usage:     "Traverse the state with given root hash and perform detailed verification",
 				ArgsUsage: "<root>",
 				Action:    traverseRawState,
-				Flags:     slices.Concat(utils.NetworkFlags, utils.DatabaseFlags),
+				Flags:     slices.Concat(utils.TraverseStateFlags, utils.NetworkFlags, utils.DatabaseFlags),
 				Description: `
-geth snapshot traverse-rawstate <state-root>
-will traverse the whole state from the given root and will abort if any referenced
-trie node or contract code is missing. This command can be used for state integrity
-verification. The default checking target is the HEAD state. It's basically identical
-to traverse-state, but the check granularity is smaller.
+geth snapshot traverse-rawstate [--account <account>] [--start <key>] [--limit <key>] <state-root>
+
+Similar to traverse-state but with more detailed verification at the trie node level.
+
+1. Traverse the whole state from the given state root:
+- --start: starting account key (64/66 chars hex) [optional]
+- --limit: ending account key (64/66 chars hex) [optional]
+
+2. Traverse a specific account's storage:
+- --account: account address (40/42 chars) or hash (64/66 chars) [required]
+- --start: starting storage key (64/66 chars hex) [optional]
+- --limit: ending storage key (64/66 chars hex) [optional]
+
+The default checking state root is the HEAD state if not specified.
+The command will abort if any referenced trie node or contract code is missing.
+This can be used for state integrity verification. The default target is HEAD state.
 
 It's also usable without snapshot enabled.
 `,
@@ -272,111 +301,549 @@ func checkDanglingStorage(ctx *cli.Context) error {
 	return snapshot.CheckDanglingStorage(db)
 }
 
-// traverseState is a helper function used for pruning verification.
-// Basically it just iterates the trie, ensure all nodes and associated
-// contract codes are present.
 func traverseState(ctx *cli.Context) error {
+	ts, err := setupTraversal(ctx)
+	if err != nil {
+		return err
+	}
+	defer ts.Close()
+
+	var (
+		counters     = &traverseCounters{start: time.Now()}
+		cctx, cancel = context.WithCancel(context.Background())
+	)
+	defer cancel()
+
+	go func() {
+		timer := time.NewTicker(time.Second * 8)
+		defer timer.Stop()
+		for {
+			select {
+			case <-timer.C:
+				log.Info("Traversing state", "accounts", counters.accounts.Load(), "slots", counters.slots.Load(), "codes", counters.codes.Load(), "elapsed", common.PrettyDuration(time.Since(counters.start)))
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	if ts.config.isAccount {
+		return ts.traverseAccount(cctx, counters, false)
+	} else {
+		return ts.traverseState(cctx, counters, false)
+	}
+}
+
+type OnStorageNodeHook func(node common.Hash, path []byte) error
+
+// traverseStorage parallelizes storage trie traversal
+func traverseStorage(ctx context.Context, storageTrie *trie.StateTrie, startKey, limitKey []byte, raw bool, hook OnStorageNodeHook) (uint64, uint64, error) {
+	var (
+		eg, cctx = errgroup.WithContext(ctx)
+		slots    atomic.Uint64
+		nodes    atomic.Uint64
+	)
+
+	for i := 0; i < 16; i++ {
+		nibble := byte(i)
+		eg.Go(func() error {
+			// Calculate this branch's natural boundaries
+			var (
+				branchStart = []byte{nibble << 4}
+				branchLimit []byte
+			)
+			if nibble < 15 {
+				branchLimit = []byte{(nibble + 1) << 4}
+			}
+
+			// Skip branches that are entirely before startKey
+			if startKey != nil && branchLimit != nil && bytes.Compare(branchLimit, startKey) <= 0 {
+				return nil
+			}
+
+			// Skip branches that are entirely after limitKey
+			if limitKey != nil && bytes.Compare(branchStart, limitKey) >= 0 {
+				return nil
+			}
+
+			// Use the more restrictive start boundary
+			if startKey != nil && bytes.Compare(branchStart, startKey) < 0 {
+				branchStart = startKey
+			}
+			if limitKey != nil && (branchLimit == nil || bytes.Compare(branchLimit, limitKey) > 0) {
+				branchLimit = limitKey
+			}
+
+			// Skip if branch range is empty
+			if branchLimit != nil && bytes.Compare(branchStart, branchLimit) >= 0 {
+				return nil
+			}
+
+			localSlots, localNodes, err := traverseStorageBranch(cctx, storageTrie, branchStart, branchLimit, raw, hook)
+			if err != nil {
+				return err
+			}
+			slots.Add(localSlots)
+			nodes.Add(localNodes)
+			return nil
+		})
+	}
+
+	if err := eg.Wait(); err != nil {
+		return 0, 0, err
+	}
+
+	return slots.Load(), nodes.Load(), nil
+}
+
+// traverseStorageBranch traverses a specific range of the storage trie using hooks
+func traverseStorageBranch(ctx context.Context, storageTrie *trie.StateTrie, startKey, limitKey []byte, raw bool, hook OnStorageNodeHook) (slots, nodes uint64, err error) {
+	nodeIter, err := storageTrie.NodeIterator(startKey)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	if raw {
+		// Raw traversal with detailed node checking
+		for nodeIter.Next(true) {
+			select {
+			case <-ctx.Done():
+				return 0, 0, ctx.Err()
+			default:
+			}
+
+			nodes++
+			if hook != nil {
+				if err := hook(nodeIter.Hash(), nodeIter.Path()); err != nil {
+					return 0, 0, err
+				}
+			}
+
+			if nodeIter.Leaf() {
+				if limitKey != nil && bytes.Compare(nodeIter.LeafKey(), limitKey) >= 0 {
+					break
+				}
+				slots++
+			}
+		}
+	} else {
+		// Simple traversal - just iterate through leaf nodes
+		storageIter := trie.NewIterator(nodeIter)
+		for storageIter.Next() {
+			select {
+			case <-ctx.Done():
+				return 0, 0, ctx.Err()
+			default:
+			}
+
+			if limitKey != nil && bytes.Compare(storageIter.Key, limitKey) >= 0 {
+				break
+			}
+			slots++
+		}
+
+		if storageIter.Err != nil {
+			return 0, 0, storageIter.Err
+		}
+	}
+
+	if err := nodeIter.Error(); err != nil {
+		return 0, 0, err
+	}
+
+	return slots, nodes, nil
+}
+
+type traverseConfig struct {
+	root      common.Hash
+	startKey  []byte
+	limitKey  []byte
+	account   common.Hash
+	isAccount bool
+}
+
+type traverseSetup struct {
+	stack   *node.Node
+	chaindb ethdb.Database
+	triedb  *triedb.Database
+	trie    *trie.StateTrie
+	config  *traverseConfig
+}
+
+type traverseCounters struct {
+	accounts atomic.Uint64
+	slots    atomic.Uint64
+	codes    atomic.Uint64
+	nodes    atomic.Uint64
+	start    time.Time
+}
+
+func setupTraversal(ctx *cli.Context) (*traverseSetup, error) {
 	stack, _ := makeConfigNode(ctx)
-	defer stack.Close()
 
 	chaindb := utils.MakeChainDatabase(ctx, stack, true)
-	defer chaindb.Close()
-
 	triedb := utils.MakeTrieDatabase(ctx, stack, chaindb, false, true, false)
-	defer triedb.Close()
 
 	headBlock := rawdb.ReadHeadBlock(chaindb)
 	if headBlock == nil {
 		log.Error("Failed to load head block")
-		return errors.New("no head block")
+		return nil, errors.New("no head block")
 	}
-	if ctx.NArg() > 1 {
-		log.Error("Too many arguments given")
-		return errors.New("too many arguments")
+
+	config, err := parseTraverseArgs(ctx)
+	if err != nil {
+		return nil, err
 	}
-	var (
-		root common.Hash
-		err  error
-	)
-	if ctx.NArg() == 1 {
-		root, err = parseRoot(ctx.Args().First())
+	if config.root == (common.Hash{}) {
+		config.root = headBlock.Root()
+	}
+
+	t, err := trie.NewStateTrie(trie.StateTrieID(config.root), triedb)
+	if err != nil {
+		log.Error("Failed to open trie", "root", config.root, "err", err)
+		return nil, err
+	}
+
+	return &traverseSetup{
+		stack:   stack,
+		chaindb: chaindb,
+		triedb:  triedb,
+		trie:    t,
+		config:  config,
+	}, nil
+}
+
+func (ts *traverseSetup) Close() {
+	ts.triedb.Close()
+	ts.chaindb.Close()
+	ts.stack.Close()
+}
+
+func (ts *traverseSetup) traverseAccount(ctx context.Context, counters *traverseCounters, raw bool) error {
+	log.Info("Start traversing storage trie", "root", ts.config.root.Hex(), "account", ts.config.account.Hex(), "startKey", common.Bytes2Hex(ts.config.startKey), "limitKey", common.Bytes2Hex(ts.config.limitKey))
+
+	acc, err := ts.trie.GetAccountByHash(ts.config.account)
+	if err != nil {
+		log.Error("Get account failed", "account", ts.config.account.Hex(), "err", err)
+		return err
+	}
+
+	if acc.Root == types.EmptyRootHash {
+		log.Info("Account has no storage")
+		return nil
+	}
+
+	id := trie.StorageTrieID(ts.config.root, ts.config.account, acc.Root)
+	storageTrie, err := trie.NewStateTrie(id, ts.triedb)
+	if err != nil {
+		log.Error("Failed to open storage trie", "root", acc.Root, "err", err)
+		return err
+	}
+
+	var hook OnStorageNodeHook
+	if raw {
+		reader, err := ts.triedb.NodeReader(ts.config.root)
 		if err != nil {
-			log.Error("Failed to resolve state root", "err", err)
-			return err
+			log.Error("State is non-existent", "root", ts.config.root)
+			return nil
 		}
-		log.Info("Start traversing the state", "root", root)
+		// Create storage hook inline - same pattern as in TraverseHooks
+		accountHash := ts.config.account
+		hook = func(node common.Hash, path []byte) error {
+			return verifyTrieNode(reader, accountHash, node, path)
+		}
+	}
+
+	slots, nodes, err := traverseStorage(ctx, storageTrie, ts.config.startKey, ts.config.limitKey, raw, hook)
+	if err != nil {
+		log.Error("Failed to traverse storage trie", "root", acc.Root, "err", err)
+		return err
+	}
+
+	counters.slots.Add(slots)
+	counters.nodes.Add(nodes)
+
+	if raw {
+		log.Info("Storage traversal complete (raw)", "nodes", counters.nodes.Load(), "slots", counters.slots.Load(), "elapsed", common.PrettyDuration(time.Since(counters.start)))
 	} else {
-		root = headBlock.Root()
-		log.Info("Start traversing the state", "root", root, "number", headBlock.NumberU64())
+		log.Info("Storage traversal complete", "slots", counters.slots.Load(), "elapsed", common.PrettyDuration(time.Since(counters.start)))
 	}
-	t, err := trie.NewStateTrie(trie.StateTrieID(root), triedb)
-	if err != nil {
-		log.Error("Failed to open trie", "root", root, "err", err)
-		return err
-	}
-	var (
-		accounts   int
-		slots      int
-		codes      int
-		lastReport time.Time
-		start      = time.Now()
-	)
-	acctIt, err := t.NodeIterator(nil)
-	if err != nil {
-		log.Error("Failed to open iterator", "root", root, "err", err)
-		return err
-	}
-	accIter := trie.NewIterator(acctIt)
-	for accIter.Next() {
-		accounts += 1
-		var acc types.StateAccount
-		if err := rlp.DecodeBytes(accIter.Value, &acc); err != nil {
-			log.Error("Invalid account encountered during traversal", "err", err)
-			return err
+	return nil
+}
+
+func (ts *traverseSetup) traverseState(ctx context.Context, counters *traverseCounters, raw bool) error {
+	log.Info("Start traversing state trie", "root", ts.config.root.Hex(), "startKey", common.Bytes2Hex(ts.config.startKey), "limitKey", common.Bytes2Hex(ts.config.limitKey))
+
+	eg, ctx := errgroup.WithContext(ctx)
+	var reader database.NodeReader
+	if raw {
+		var err error
+		reader, err = ts.triedb.NodeReader(ts.config.root)
+		if err != nil {
+			log.Error("State is non-existent", "root", ts.config.root)
+			return nil
 		}
+	}
+
+	for i := 0; i < 16; i++ {
+		nibble := byte(i)
+		eg.Go(func() error {
+			var (
+				startKey = []byte{nibble << 4}
+				limitKey []byte
+			)
+			if nibble < 15 {
+				limitKey = []byte{(nibble + 1) << 4}
+			}
+
+			if ts.config != nil {
+				// Skip branches that are entirely before startKey
+				if ts.config.startKey != nil && limitKey != nil && bytes.Compare(limitKey, ts.config.startKey) <= 0 {
+					return nil
+				}
+
+				// Skip branches that are entirely after limitKey
+				if ts.config.limitKey != nil && bytes.Compare(startKey, ts.config.limitKey) >= 0 {
+					return nil
+				}
+
+				if ts.config.startKey != nil && bytes.Compare(startKey, ts.config.startKey) < 0 {
+					startKey = ts.config.startKey
+				}
+				if ts.config.limitKey != nil && (limitKey == nil || bytes.Compare(limitKey, ts.config.limitKey) > 0) {
+					limitKey = ts.config.limitKey
+				}
+			}
+
+			if limitKey != nil && bytes.Compare(startKey, limitKey) >= 0 {
+				return nil
+			}
+
+			var hooks *TraverseHooks
+			if raw {
+				hooks = ts.createRawHooks(reader)
+			} else {
+				hooks = ts.createSimpleHooks()
+			}
+
+			return ts.traverseStateBranch(ctx, startKey, limitKey, raw, counters, hooks)
+		})
+	}
+
+	if err := eg.Wait(); err != nil {
+		return err
+	}
+
+	log.Info("State traversal complete", "accounts", counters.accounts.Load(), "slots", counters.slots.Load(), "codes", counters.codes.Load(), "elapsed", common.PrettyDuration(time.Since(counters.start)))
+	return nil
+}
+
+func parseTraverseArgs(ctx *cli.Context) (*traverseConfig, error) {
+	if ctx.NArg() > 1 {
+		return nil, errors.New("too many arguments, only <root> is required")
+	}
+
+	config := &traverseConfig{}
+	var err error
+
+	if ctx.NArg() == 1 {
+		config.root, err = parseRoot(ctx.Args().First())
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if accountFlag := ctx.String("account"); accountFlag != "" {
+		config.isAccount = true
+		switch len(accountFlag) {
+		case 40, 42:
+			config.account = crypto.Keccak256Hash(common.HexToAddress(accountFlag).Bytes())
+		case 64, 66:
+			config.account = common.HexToHash(accountFlag)
+		default:
+			return nil, errors.New("account must be 40/42 chars for address or 64/66 chars for hash")
+		}
+	}
+
+	if startFlag := ctx.String("start"); startFlag != "" {
+		if len(startFlag) == 64 || len(startFlag) == 66 {
+			config.startKey = common.HexToHash(startFlag).Bytes()
+		} else {
+			return nil, errors.New("start key must be 64/66 chars hex")
+		}
+	}
+
+	if limitFlag := ctx.String("limit"); limitFlag != "" {
+		if len(limitFlag) == 64 || len(limitFlag) == 66 {
+			config.limitKey = common.HexToHash(limitFlag).Bytes()
+		} else {
+			return nil, errors.New("limit key must be 64/66 chars hex")
+		}
+	}
+
+	return config, nil
+}
+
+// TraverseHooks defines the hooks for different traversal modes
+type TraverseHooks struct {
+	// Called for each trie account node (only for raw mode)
+	OnAccountNode func(node common.Hash, path []byte) error
+	// Called for each storage trie - now self-contained with all needed context
+	OnStorageTrie func(ctx context.Context, storageTrie *trie.StateTrie, accountHash common.Hash) (slots, nodes uint64, err error)
+}
+
+// createSimpleHooks creates hooks for simple traversal mode
+func (ts *traverseSetup) createSimpleHooks() *TraverseHooks {
+	return &TraverseHooks{
+		OnStorageTrie: func(ctx context.Context, storageTrie *trie.StateTrie, accountHash common.Hash) (slots, nodes uint64, err error) {
+			return traverseStorage(ctx, storageTrie, nil, nil, false, nil)
+		},
+	}
+}
+
+func verifyTrieNode(reader database.NodeReader, accountHash common.Hash, node common.Hash, path []byte) error {
+	if node != (common.Hash{}) {
+		blob, _ := reader.Node(accountHash, path, node)
+		if len(blob) == 0 {
+			log.Error("Missing trie node", "hash", node)
+			return errors.New("missing node")
+		}
+		if !bytes.Equal(crypto.Keccak256(blob), node.Bytes()) {
+			log.Error("Invalid trie node", "hash", node.Hex(), "value", blob)
+			return errors.New("invalid node")
+		}
+	}
+	return nil
+}
+
+// createRawHooks creates hooks for raw traversal mode with detailed verification
+func (ts *traverseSetup) createRawHooks(reader database.NodeReader) *TraverseHooks {
+	return &TraverseHooks{
+		OnAccountNode: func(node common.Hash, path []byte) error {
+			return verifyTrieNode(reader, common.Hash{}, node, path)
+		},
+
+		OnStorageTrie: func(ctx context.Context, storageTrie *trie.StateTrie, accountHash common.Hash) (slots, nodes uint64, err error) {
+			hook := func(node common.Hash, path []byte) error {
+				return verifyTrieNode(reader, accountHash, node, path)
+			}
+			return traverseStorage(ctx, storageTrie, nil, nil, true, hook)
+		},
+	}
+}
+
+// traverseStateBranch provides common branch traversal logic using hooks
+func (ts *traverseSetup) traverseStateBranch(ctx context.Context, startKey, limitKey []byte, raw bool, counters *traverseCounters, hooks *TraverseHooks) error {
+	accIter, err := ts.trie.NodeIterator(startKey)
+	if err != nil {
+		return err
+	}
+
+	processAccount := func(accKey []byte, acc *types.StateAccount) error {
+		accountHash := common.BytesToHash(accKey)
+
+		// Process storage if present
 		if acc.Root != types.EmptyRootHash {
-			id := trie.StorageTrieID(root, common.BytesToHash(accIter.Key), acc.Root)
-			storageTrie, err := trie.NewStateTrie(id, triedb)
+			id := trie.StorageTrieID(ts.config.root, accountHash, acc.Root)
+			storageTrie, err := trie.NewStateTrie(id, ts.triedb)
 			if err != nil {
 				log.Error("Failed to open storage trie", "root", acc.Root, "err", err)
 				return err
 			}
-			storageIt, err := storageTrie.NodeIterator(nil)
+
+			slots, nodes, err := hooks.OnStorageTrie(ctx, storageTrie, accountHash)
 			if err != nil {
-				log.Error("Failed to open storage iterator", "root", acc.Root, "err", err)
 				return err
 			}
-			storageIter := trie.NewIterator(storageIt)
-			for storageIter.Next() {
-				slots += 1
-
-				if time.Since(lastReport) > time.Second*8 {
-					log.Info("Traversing state", "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
-					lastReport = time.Now()
-				}
-			}
-			if storageIter.Err != nil {
-				log.Error("Failed to traverse storage trie", "root", acc.Root, "err", storageIter.Err)
-				return storageIter.Err
-			}
+			counters.slots.Add(slots)
+			counters.nodes.Add(nodes)
 		}
+
+		// Process code if present
 		if !bytes.Equal(acc.CodeHash, types.EmptyCodeHash.Bytes()) {
-			if !rawdb.HasCode(chaindb, common.BytesToHash(acc.CodeHash)) {
-				log.Error("Code is missing", "hash", common.BytesToHash(acc.CodeHash))
+			if codeHash := common.BytesToHash(acc.CodeHash); !rawdb.HasCode(ts.chaindb, codeHash) {
+				log.Error("Code is missing", "hash", codeHash)
 				return errors.New("missing code")
 			}
-			codes += 1
+			counters.codes.Add(1)
 		}
-		if time.Since(lastReport) > time.Second*8 {
-			log.Info("Traversing state", "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
-			lastReport = time.Now()
+		return nil
+	}
+
+	if raw {
+		// Raw traversal with detailed node checking
+		for accIter.Next(true) {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+			}
+
+			counters.nodes.Add(1)
+			if hooks != nil && hooks.OnAccountNode != nil {
+				if err := hooks.OnAccountNode(accIter.Hash(), accIter.Path()); err != nil {
+					return err
+				}
+			}
+
+			// If it's a leaf node, process the account
+			if accIter.Leaf() {
+				if limitKey != nil && bytes.Compare(accIter.LeafKey(), limitKey) >= 0 {
+					break
+				}
+				counters.accounts.Add(1)
+
+				var acc types.StateAccount
+				if err := rlp.DecodeBytes(accIter.LeafBlob(), &acc); err != nil {
+					log.Error("Invalid account encountered during traversal", "err", err)
+					return errors.New("invalid account")
+				}
+
+				if err := processAccount(accIter.LeafKey(), &acc); err != nil {
+					return err
+				}
+			}
+		}
+	} else {
+		// Simple traversal - just iterate through leaf nodes
+		acctIt, err := ts.trie.NodeIterator(startKey)
+		if err != nil {
+			return err
+		}
+
+		accIter := trie.NewIterator(acctIt)
+		for accIter.Next() {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+			}
+
+			if limitKey != nil && bytes.Compare(accIter.Key, limitKey) >= 0 {
+				break
+			}
+
+			counters.accounts.Add(1)
+
+			var acc types.StateAccount
+			if err := rlp.DecodeBytes(accIter.Value, &acc); err != nil {
+				log.Error("Invalid account encountered during traversal", "err", err)
+				return err
+			}
+			if err := processAccount(accIter.Key, &acc); err != nil {
+				return err
+			}
+		}
+
+		if accIter.Err != nil {
+			return accIter.Err
 		}
 	}
-	if accIter.Err != nil {
-		log.Error("Failed to traverse state trie", "root", root, "err", accIter.Err)
-		return accIter.Err
+
+	if accIter.Error() != nil {
+		return accIter.Error()
 	}
-	log.Info("State is complete", "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
+
 	return nil
 }
 
@@ -385,158 +852,36 @@ func traverseState(ctx *cli.Context) error {
 // contract codes are present. It's basically identical to traverseState
 // but it will check each trie node.
 func traverseRawState(ctx *cli.Context) error {
-	stack, _ := makeConfigNode(ctx)
-	defer stack.Close()
-
-	chaindb := utils.MakeChainDatabase(ctx, stack, true)
-	defer chaindb.Close()
-
-	triedb := utils.MakeTrieDatabase(ctx, stack, chaindb, false, true, false)
-	defer triedb.Close()
-
-	headBlock := rawdb.ReadHeadBlock(chaindb)
-	if headBlock == nil {
-		log.Error("Failed to load head block")
-		return errors.New("no head block")
-	}
-	if ctx.NArg() > 1 {
-		log.Error("Too many arguments given")
-		return errors.New("too many arguments")
+	ts, err := setupTraversal(ctx)
+	if err != nil {
+		return err
 	}
+	defer ts.Close()
+
 	var (
-		root common.Hash
-		err  error
+		counters     = &traverseCounters{start: time.Now()}
+		cctx, cancel = context.WithCancel(context.Background())
 	)
-	if ctx.NArg() == 1 {
-		root, err = parseRoot(ctx.Args().First())
-		if err != nil {
-			log.Error("Failed to resolve state root", "err", err)
-			return err
+	defer cancel()
+
+	go func() {
+		timer := time.NewTicker(time.Second * 8)
+		defer timer.Stop()
+		for {
+			select {
+			case <-timer.C:
+				log.Info("Traversing rawstate", "nodes", counters.nodes.Load(), "accounts", counters.accounts.Load(), "slots", counters.slots.Load(), "codes", counters.codes.Load(), "elapsed", common.PrettyDuration(time.Since(counters.start)))
+			case <-cctx.Done():
+				return
+			}
 		}
-		log.Info("Start traversing the state", "root", root)
+	}()
+
+	if ts.config.isAccount {
+		return ts.traverseAccount(cctx, counters, true)
 	} else {
-		root = headBlock.Root()
-		log.Info("Start traversing the state", "root", root, "number", headBlock.NumberU64())
+		return ts.traverseState(cctx, counters, true)
 	}
-	t, err := trie.NewStateTrie(trie.StateTrieID(root), triedb)
-	if err != nil {
-		log.Error("Failed to open trie", "root", root, "err", err)
-		return err
-	}
-	var (
-		nodes      int
-		accounts   int
-		slots      int
-		codes      int
-		lastReport time.Time
-		start      = time.Now()
-		hasher     = crypto.NewKeccakState()
-		got        = make([]byte, 32)
-	)
-	accIter, err := t.NodeIterator(nil)
-	if err != nil {
-		log.Error("Failed to open iterator", "root", root, "err", err)
-		return err
-	}
-	reader, err := triedb.NodeReader(root)
-	if err != nil {
-		log.Error("State is non-existent", "root", root)
-		return nil
-	}
-	for accIter.Next(true) {
-		nodes += 1
-		node := accIter.Hash()
-
-		// Check the present for non-empty hash node(embedded node doesn't
-		// have their own hash).
-		if node != (common.Hash{}) {
-			blob, _ := reader.Node(common.Hash{}, accIter.Path(), node)
-			if len(blob) == 0 {
-				log.Error("Missing trie node(account)", "hash", node)
-				return errors.New("missing account")
-			}
-			hasher.Reset()
-			hasher.Write(blob)
-			hasher.Read(got)
-			if !bytes.Equal(got, node.Bytes()) {
-				log.Error("Invalid trie node(account)", "hash", node.Hex(), "value", blob)
-				return errors.New("invalid account node")
-			}
-		}
-		// If it's a leaf node, yes we are touching an account,
-		// dig into the storage trie further.
-		if accIter.Leaf() {
-			accounts += 1
-			var acc types.StateAccount
-			if err := rlp.DecodeBytes(accIter.LeafBlob(), &acc); err != nil {
-				log.Error("Invalid account encountered during traversal", "err", err)
-				return errors.New("invalid account")
-			}
-			if acc.Root != types.EmptyRootHash {
-				id := trie.StorageTrieID(root, common.BytesToHash(accIter.LeafKey()), acc.Root)
-				storageTrie, err := trie.NewStateTrie(id, triedb)
-				if err != nil {
-					log.Error("Failed to open storage trie", "root", acc.Root, "err", err)
-					return errors.New("missing storage trie")
-				}
-				storageIter, err := storageTrie.NodeIterator(nil)
-				if err != nil {
-					log.Error("Failed to open storage iterator", "root", acc.Root, "err", err)
-					return err
-				}
-				for storageIter.Next(true) {
-					nodes += 1
-					node := storageIter.Hash()
-
-					// Check the presence for non-empty hash node(embedded node doesn't
-					// have their own hash).
-					if node != (common.Hash{}) {
-						blob, _ := reader.Node(common.BytesToHash(accIter.LeafKey()), storageIter.Path(), node)
-						if len(blob) == 0 {
-							log.Error("Missing trie node(storage)", "hash", node)
-							return errors.New("missing storage")
-						}
-						hasher.Reset()
-						hasher.Write(blob)
-						hasher.Read(got)
-						if !bytes.Equal(got, node.Bytes()) {
-							log.Error("Invalid trie node(storage)", "hash", node.Hex(), "value", blob)
-							return errors.New("invalid storage node")
-						}
-					}
-					// Bump the counter if it's leaf node.
-					if storageIter.Leaf() {
-						slots += 1
-					}
-					if time.Since(lastReport) > time.Second*8 {
-						log.Info("Traversing state", "nodes", nodes, "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
-						lastReport = time.Now()
-					}
-				}
-				if storageIter.Error() != nil {
-					log.Error("Failed to traverse storage trie", "root", acc.Root, "err", storageIter.Error())
-					return storageIter.Error()
-				}
-			}
-			if !bytes.Equal(acc.CodeHash, types.EmptyCodeHash.Bytes()) {
-				if !rawdb.HasCode(chaindb, common.BytesToHash(acc.CodeHash)) {
-					log.Error("Code is missing", "account", common.BytesToHash(accIter.LeafKey()))
-					return errors.New("missing code")
-				}
-				codes += 1
-			}
-			if time.Since(lastReport) > time.Second*8 {
-				log.Info("Traversing state", "nodes", nodes, "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
-				lastReport = time.Now()
-			}
-		}
-	}
-	if accIter.Error() != nil {
-		log.Error("Failed to traverse state trie", "root", root, "err", accIter.Error())
-		return accIter.Error()
-	}
-	log.Info("State is complete", "nodes", nodes, "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
-	return nil
 }
 
 func parseRoot(input string) (common.Hash, error) {

cmd/utils/flags.go

@@ -1130,6 +1130,22 @@ var (
 		StateSchemeFlag,
 		HttpHeaderFlag,
 	}
+
+	// TraverseStateFlags is the flag group of all state traversal flags.
+	TraverseStateFlags = []cli.Flag{
+		&cli.StringFlag{
+			Name:  "account",
+			Usage: "Account address or hash to traverse storage for (enables account mode)",
+		},
+		&cli.StringFlag{
+			Name:  "start",
+			Usage: "Starting key (account key for state mode, storage key for account mode)",
+		},
+		&cli.StringFlag{
+			Name:  "limit",
+			Usage: "Ending key (account key for state mode, storage key for account mode)",
+		},
+	}
 )
 
 // default account to prefund when running Geth in dev mode