cmd/geth, triedb: add generate-trie bench. fix iterator key copy

cmd/geth/snapshot.go

@@ -23,10 +23,15 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"os/signal"
+	"path/filepath"
 	"slices"
 	"sort"
+	"strings"
+	"syscall"
 	"time"
 
+	pebbleimpl "github.com/cockroachdb/pebble"
 	"github.com/ethereum/go-ethereum/cmd/utils"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core"
@@ -36,6 +41,7 @@ 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/pebble"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
@@ -80,6 +86,33 @@ geth snapshot verify-state <state-root>
 will traverse the whole accounts and storages set based on the specified
 snapshot and recalculate the root hash of state for verification.
 In other words, this command does the snapshot to trie conversion.
+`,
+			},
+			{
+				Name:      "generate-trie",
+				Usage:     "Benchmark triedb.GenerateTrie against a hardlinked checkpoint of the chaindata",
+				ArgsUsage: "[<root>]",
+				Action:    benchGenerateTrie,
+				Flags: slices.Concat(utils.NetworkFlags, utils.DatabaseFlags, []cli.Flag{
+					&cli.StringFlag{
+						Name:  "checkpoint",
+						Usage: "Directory for the pebble checkpoint (default: <chaindata-parent>/.gentrie-bench-<ts>)",
+					},
+					&cli.BoolFlag{
+						Name:  "keep",
+						Usage: "Keep the checkpoint directory after the run (debugging)",
+					},
+				}),
+				Description: `
+geth snapshot generate-trie [<root>]
+
+Takes a pebble checkpoint of the chaindata (hardlinked SST files, near-zero
+disk usage and near-instant) and runs triedb.GenerateTrie against the
+checkpoint. The source datadir is opened read-only for the checkpoint and
+never written to. The checkpoint is removed on exit unless --keep is set,
+including on Ctrl-C.
+
+If <root> is not given, the head block's root is used.
 `,
 			},
 			{
@@ -289,6 +322,151 @@ func verifyState(ctx *cli.Context) error {
 	}
 }
 
+// benchGenerateTrie runs triedb.GenerateTrie against a hardlinked checkpoint
+// of the chaindata so the source datadir is never written to.
+func benchGenerateTrie(ctx *cli.Context) error {
+	stack, _ := makeConfigNode(ctx)
+	defer stack.Close()
+
+	// Resolve source chaindata path (handles network-specific subdirs).
+	srcDir := stack.ResolvePath("chaindata")
+	if fi, err := os.Stat(srcDir); err != nil {
+		return fmt.Errorf("chaindata not found at %s: %w", srcDir, err)
+	} else if !fi.IsDir() {
+		return fmt.Errorf("%s is not a directory", srcDir)
+	}
+
+	// Default to snapshot root, not head: that's what GenerateTrie actually
+	// reconstructs from flat state. On a fully-synced node they match.
+	var root common.Hash
+	if ctx.NArg() == 1 {
+		r, err := parseRoot(ctx.Args().First())
+		if err != nil {
+			return fmt.Errorf("parse root: %w", err)
+		}
+		root = r
+	} else {
+		chaindb := utils.MakeChainDatabase(ctx, stack, true)
+		snapRoot := rawdb.ReadSnapshotRoot(chaindb)
+		head := rawdb.ReadHeadBlock(chaindb)
+		chaindb.Close()
+		switch {
+		case snapRoot != (common.Hash{}):
+			root = snapRoot
+			log.Info("using snapshot root", "root", root)
+		case head != nil:
+			root = head.Root()
+			log.Info("using head block root", "number", head.Number(), "root", root)
+		default:
+			return errors.New("no snapshot or head block found; pass <root> explicitly")
+		}
+	}
+
+	// Default checkpoint sits next to chaindata so hardlinks work.
+	ckpt := ctx.String("checkpoint")
+	if ckpt == "" {
+		ts := time.Now().Format("20060102-150405")
+		ckpt = filepath.Join(filepath.Dir(srcDir), fmt.Sprintf(".gentrie-bench-%s", ts))
+	}
+	if _, err := os.Stat(ckpt); err == nil {
+		return fmt.Errorf("checkpoint dir %s already exists; remove it or pass --checkpoint to a fresh path", ckpt)
+	}
+
+	log.Info("creating pebble checkpoint", "src", srcDir, "dst", ckpt)
+	cpStart := time.Now()
+	if err := makeCheckpoint(srcDir, ckpt); err != nil {
+		return fmt.Errorf("checkpoint failed: %w", err)
+	}
+	log.Info("checkpoint created", "elapsed", time.Since(cpStart))
+
+	// Clean up the checkpoint on exit, including Ctrl-C.
+	keep := ctx.Bool("keep")
+	cleanup := func() {
+		if keep {
+			log.Info("keeping checkpoint", "path", ckpt)
+			return
+		}
+		log.Info("removing checkpoint", "path", ckpt)
+		if err := os.RemoveAll(ckpt); err != nil {
+			log.Error("failed to remove checkpoint", "err", err)
+		}
+	}
+	defer cleanup()
+
+	cancelCh := make(chan struct{})
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
+	defer signal.Stop(sigCh)
+	go func() {
+		<-sigCh
+		log.Warn("interrupt received; cancelling GenerateTrie")
+		close(cancelCh)
+	}()
+
+	// Open the checkpoint writable. Reuse source ancient. Checkpoint only
+	// hardlinks the pebble SSTs (not the freezer), and GenerateTrie never
+	// writes to ancient, so sharing it is safe.
+	srcAncient := stack.ResolveAncient("chaindata", "")
+	kv, err := pebble.New(ckpt, 4096, 1024, "gentrie-bench", false)
+	if err != nil {
+		return fmt.Errorf("open checkpoint: %w", err)
+	}
+	chaindb, err := rawdb.Open(kv, rawdb.OpenOptions{
+		Ancient:          srcAncient,
+		MetricsNamespace: "gentrie-bench",
+	})
+	if err != nil {
+		kv.Close()
+		return fmt.Errorf("rawdb.Open checkpoint: %w", err)
+	}
+	defer chaindb.Close()
+
+	// Pick up the trie scheme already in use (path or hash).
+	triedbInst := utils.MakeTrieDatabase(ctx, stack, chaindb, false, true, false)
+	scheme := triedbInst.Scheme()
+	triedbInst.Close()
+
+	log.Info("running GenerateTrie", "scheme", scheme, "root", root)
+	runStart := time.Now()
+	err = triedb.GenerateTrie(chaindb, scheme, root, cancelCh)
+	elapsed := time.Since(runStart)
+	if err != nil {
+		// On a mid-snap-sync datadir the reconstructed root won't match the
+		// expected one. Treat that as a warning so the benchmark still
+		// reports wall time. Real errors (iterator, write failures) propagate.
+		if strings.Contains(err.Error(), "state root mismatch") {
+			log.Warn("root mismatch (expected on partial snapshot)", "err", err)
+		} else {
+			log.Error("GenerateTrie failed", "elapsed", elapsed, "err", err)
+			return err
+		}
+	}
+
+	fmt.Printf("\n=== generate-trie benchmark ===\n")
+	fmt.Printf("source:     %s (untouched)\n", srcDir)
+	fmt.Printf("checkpoint: %s\n", ckpt)
+	fmt.Printf("scheme:     %s\n", scheme)
+	fmt.Printf("root:       %s\n", root.Hex())
+	fmt.Printf("wall time:  %s\n", elapsed)
+	return nil
+}
+
+// makeCheckpoint opens srcDir as a pebble database and writes a hardlinked
+// checkpoint to dstDir. Source is closed on return.
+//
+// Opens read-write so pebble can finalize its startup (WAL replay, fresh
+// OPTIONS file) before checkpointing. Read-only mode skips that step, and
+// Checkpoint then fails trying to hardlink the missing OPTIONS file. The
+// read-write open does no more than a normal geth startup would.
+func makeCheckpoint(srcDir, dstDir string) error {
+	db, err := pebbleimpl.Open(srcDir, &pebbleimpl.Options{})
+	if err != nil {
+		return fmt.Errorf("open source pebble: %w", err)
+	}
+	defer db.Close()
+	return db.Checkpoint(dstDir)
+}
+
 // checkDanglingStorage iterates the snap storage data, and verifies that all
 // storage also has corresponding account data.
 func checkDanglingStorage(ctx *cli.Context) error {

triedb/generate.go

@@ -88,7 +88,9 @@ func reopenFlatIterator(db ethdb.Database, old *internal.HoldableIterator, prefi
 		old.Release()
 		return internal.NewHoldableIterator(memorydb.New().NewIterator(nil, nil))
 	}
-	next := old.Key()
+	// pebble's Key() slice is invalidated by Release. Copy first so the new
+	// iterator's lower bound isn't seeded from freed memory.
+	next := common.CopyBytes(old.Key())
 	old.Release()
 	return openFlatIterator(db, prefix, next[len(prefix):], suffixLen)
 }