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

cmd/geth/chaincmd.go

@@ -165,6 +165,17 @@ It's deprecated, please use "geth db export" instead.
 		}, utils.DatabaseFlags),
 		Description: `
 This command dumps out the state for a given block (or latest, if none provided).
+`,
+	}
+	exportSnapshotPreimagesCommand = &cli.Command{
+		Action:    exportSnapshotPreimages,
+		Name:      "export-snapshot-preimages",
+		Usage:     "Export the preimage in snapshot enumeration order",
+		ArgsUsage: "<dumpfile> [<root>]",
+		Flags:     utils.DatabaseFlags,
+		Description: `
+The export-snapshot-preimages command exports hash preimages to a flat file, in exactly
+the expected order for the overlay tree migration.
 `,
 	}
 )
@@ -406,6 +417,33 @@ func exportPreimages(ctx *cli.Context) error {
 	return nil
 }
 
+// exportSnapshotPreimages dumps the preimage data to a flat file.
+func exportSnapshotPreimages(ctx *cli.Context) error {
+	if ctx.Args().Len() < 1 {
+		utils.Fatalf("This command requires an argument.")
+	}
+	stack, _ := makeConfigNode(ctx)
+	defer stack.Close()
+
+	chain, _ := utils.MakeChain(ctx, stack, true)
+
+	var root common.Hash
+	if ctx.Args().Len() > 1 {
+		rootBytes := common.FromHex(ctx.Args().Get(1))
+		if len(rootBytes) != common.HashLength {
+			return fmt.Errorf("invalid root hash length")
+		}
+		root = common.BytesToHash(rootBytes)
+	}
+
+	start := time.Now()
+	if err := utils.ExportSnapshotPreimages(chain, ctx.Args().First(), root); 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

@@ -209,6 +209,7 @@ func init() {
 		exportCommand,
 		importPreimagesCommand,
 		exportPreimagesCommand,
+		exportSnapshotPreimagesCommand,
 		removedbCommand,
 		dumpCommand,
 		dumpGenesisCommand,

cmd/utils/cmd.go

@@ -374,6 +374,85 @@ 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(chain *core.BlockChain, 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()
+
+	writer := bufio.NewWriter(fh)
+	defer writer.Flush()
+
+	statedb, err := chain.State()
+	if err != nil {
+		return fmt.Errorf("failed to open statedb: %w", err)
+	}
+
+	if root == (common.Hash{}) {
+		root = chain.CurrentBlock().Root
+	}
+
+	type hashAndPreimageSize struct {
+		Hash common.Hash
+		Size int
+	}
+	hashCh := make(chan hashAndPreimageSize)
+
+	go func() {
+		defer close(hashCh)
+		accIt, err := chain.Snapshots().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 := chain.Snapshots().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(statedb.Database().DiskDB(), item.Hash)
+		if len(preimage) != item.Size {
+			return fmt.Errorf("invalid preimage size")
+		}
+		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