cmd/{geth,utils}: add cmd to export preimages in snap enumeration order

cmd/geth/chaincmd.go

@@ -120,34 +120,6 @@ Optional second and third arguments control the first and
 last block to write. In this mode, the file will be appended
 if already existing. If the file ends with .gz, the output will
 be gzipped.`,
-	}
-	importPreimagesCommand = &cli.Command{
-		Action:    importPreimages,
-		Name:      "import-preimages",
-		Usage:     "Import the preimage database from an RLP stream",
-		ArgsUsage: "<datafile>",
-		Flags: flags.Merge([]cli.Flag{
-			utils.CacheFlag,
-			utils.SyncModeFlag,
-		}, utils.DatabaseFlags),
-		Description: `
-The import-preimages command imports hash preimages from an RLP encoded stream.
-It's deprecated, please use "geth db import" instead.
-`,
-	}
-	exportPreimagesCommand = &cli.Command{
-		Action:    exportPreimages,
-		Name:      "export-preimages",
-		Usage:     "Export the preimage database into an RLP stream",
-		ArgsUsage: "<dumpfile>",
-		Flags: flags.Merge([]cli.Flag{
-			utils.CacheFlag,
-			utils.SyncModeFlag,
-		}, utils.DatabaseFlags),
-		Description: `
-The export-preimages command exports hash preimages to an RLP encoded stream.
-It's deprecated, please use "geth db export" instead.
-`,
 	}
 	dumpCommand = &cli.Command{
 		Action:    dump,
@@ -367,45 +339,6 @@ func exportChain(ctx *cli.Context) error {
 	return nil
 }
 
-// importPreimages imports preimage data from the specified file.
-func importPreimages(ctx *cli.Context) error {
-	if ctx.Args().Len() < 1 {
-		utils.Fatalf("This command requires an argument.")
-	}
-
-	stack, _ := makeConfigNode(ctx)
-	defer stack.Close()
-
-	db := utils.MakeChainDatabase(ctx, stack, false)
-	defer db.Close()
-	start := time.Now()
-
-	if err := utils.ImportPreimages(db, ctx.Args().First()); err != nil {
-		utils.Fatalf("Import error: %v\n", err)
-	}
-	fmt.Printf("Import done in %v\n", time.Since(start))
-	return nil
-}
-
-// exportPreimages dumps the preimage data to specified json file in streaming way.
-func exportPreimages(ctx *cli.Context) error {
-	if ctx.Args().Len() < 1 {
-		utils.Fatalf("This command requires an argument.")
-	}
-	stack, _ := makeConfigNode(ctx)
-	defer stack.Close()
-
-	db := utils.MakeChainDatabase(ctx, stack, true)
-	defer db.Close()
-	start := time.Now()
-
-	if err := utils.ExportPreimages(db, ctx.Args().First()); err != nil {
-		utils.Fatalf("Export error: %v\n", err)
-	}
-	fmt.Printf("Export done in %v\n", time.Since(start))
-	return nil
-}
-
 func parseDumpConfig(ctx *cli.Context, stack *node.Node) (*state.DumpConfig, ethdb.Database, common.Hash, error) {
 	db := utils.MakeChainDatabase(ctx, stack, true)
 	defer db.Close()

cmd/geth/main.go

@@ -207,8 +207,6 @@ func init() {
 		initCommand,
 		importCommand,
 		exportCommand,
-		importPreimagesCommand,
-		exportPreimagesCommand,
 		removedbCommand,
 		dumpCommand,
 		dumpGenesisCommand,

cmd/geth/snapshot.go

@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"os"
 	"time"
 
@@ -147,6 +148,17 @@ as the backend data source, making this command a lot faster.
 
 The argument is interpreted as block number or hash. If none is provided, the latest
 block is used.
+`,
+			},
+			{
+				Action:    snapshotExportPreimages,
+				Name:      "export-preimages",
+				Usage:     "Export the preimage in snapshot enumeration order",
+				ArgsUsage: "<dumpfile> [<root>]",
+				Flags:     utils.DatabaseFlags,
+				Description: `
+The export-preimages command exports hash preimages to a flat file, in exactly
+the expected order for the overlay tree migration.
 `,
 			},
 		},
@@ -604,6 +616,55 @@ func dumpState(ctx *cli.Context) error {
 	return nil
 }
 
+// snapshotExportPreimages dumps the preimage data to a flat file.
+func snapshotExportPreimages(ctx *cli.Context) error {
+	if ctx.NArg() < 1 {
+		utils.Fatalf("This command requires an argument.")
+	}
+	stack, _ := makeConfigNode(ctx)
+	defer stack.Close()
+
+	chaindb := utils.MakeChainDatabase(ctx, stack, true)
+	defer chaindb.Close()
+
+	triedb := utils.MakeTrieDatabase(ctx, chaindb, false, true)
+	defer triedb.Close()
+
+	var root common.Hash
+	if ctx.NArg() > 1 {
+		rootBytes := common.FromHex(ctx.Args().Get(1))
+		if len(rootBytes) != common.HashLength {
+			return fmt.Errorf("invalid hash: %s", ctx.Args().Get(1))
+		}
+		root = common.BytesToHash(rootBytes)
+	} else {
+		headBlock := rawdb.ReadHeadBlock(chaindb)
+		if headBlock == nil {
+			log.Error("Failed to load head block")
+			return errors.New("no head block")
+		}
+		root = headBlock.Root()
+	}
+
+	snapConfig := snapshot.Config{
+		CacheSize:  256,
+		Recovery:   false,
+		NoBuild:    true,
+		AsyncBuild: false,
+	}
+	snaptree, err := snapshot.New(snapConfig, chaindb, triedb, root)
+	if err != nil {
+		return err
+	}
+
+	start := time.Now()
+	if err := utils.ExportSnapshotPreimages(chaindb, snaptree, ctx.Args().First(), root); err != nil {
+		utils.Fatalf("Export error: %v\n", err)
+	}
+	log.Info("Export done", "duration", time.Since(start))
+	return nil
+}
+
 // checkAccount iterates the snap data layers, and looks up the given account
 // across all layers.
 func checkAccount(ctx *cli.Context) error {

cmd/utils/cmd.go

@@ -33,6 +33,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core"
 	"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/eth/ethconfig"
@@ -374,6 +375,88 @@ func ExportPreimages(db ethdb.Database, fn string) error {
 	return nil
 }
 
+// ExportSnapshotPreimages exports the preimages corresponding to the enumeration of
+// the snapshot for a given root.
+func ExportSnapshotPreimages(chaindb ethdb.Database, snaptree *snapshot.Tree, fn string, root common.Hash) error {
+	log.Info("Exporting preimages", "file", fn)
+
+	fh, err := os.OpenFile(fn, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
+	if err != nil {
+		return err
+	}
+	defer fh.Close()
+
+	var writer io.Writer = fh
+
+	if strings.HasSuffix(fn, ".gz") {
+		gz := gzip.NewWriter(writer)
+		defer gz.Close()
+		writer = gz
+	}
+
+	buf := bufio.NewWriter(writer)
+	defer buf.Flush()
+	writer = buf
+
+	type hashAndPreimageSize struct {
+		Hash common.Hash
+		Size int
+	}
+	hashCh := make(chan hashAndPreimageSize)
+
+	go func() {
+		defer close(hashCh)
+		accIt, err := snaptree.AccountIterator(root, common.Hash{})
+		if err != nil {
+			log.Error("Failed to create account iterator", "error", err)
+			return
+		}
+		defer accIt.Release()
+
+		count := 0
+		for accIt.Next() {
+			acc, err := types.FullAccount(accIt.Account())
+			if err != nil {
+				log.Error("Failed to get full account", "error", err)
+				return
+			}
+			hashCh <- hashAndPreimageSize{Hash: accIt.Hash(), Size: 20}
+
+			if acc.Root != (common.Hash{}) && acc.Root != types.EmptyRootHash {
+				stIt, err := snaptree.StorageIterator(root, accIt.Hash(), common.Hash{})
+				if err != nil {
+					log.Error("Failed to create storage iterator", "error", err)
+					return
+				}
+				for stIt.Next() {
+					hashCh <- hashAndPreimageSize{Hash: stIt.Hash(), Size: 32}
+				}
+				stIt.Release()
+			}
+			count++
+			if count%100000 == 0 {
+				log.Info("Last exported account", "account", accIt.Hash())
+			}
+		}
+	}()
+
+	for item := range hashCh {
+		preimage := rawdb.ReadPreimage(chaindb, item.Hash)
+		if len(preimage) == 0 {
+			return fmt.Errorf("missing preimage for %v", item.Hash)
+		}
+		if len(preimage) != item.Size {
+			return fmt.Errorf("invalid preimage size, have %d", len(preimage))
+		}
+		if _, err := writer.Write(preimage); err != nil {
+			return fmt.Errorf("failed to write preimage: %w", err)
+		}
+	}
+
+	log.Info("Exported preimages", "file", fn)
+	return nil
+}
+
 // exportHeader is used in the export/import flow. When we do an export,
 // the first element we output is the exportHeader.
 // Whenever a backwards-incompatible change is made, the Version header