From b541aa529d62e6b130c22d07d5d40199351e7532 Mon Sep 17 00:00:00 2001
From: Gary Rong <garyrong0905@gmail.com>
Date: Mon, 6 Nov 2023 21:02:00 +0800
Subject: [PATCH] cmd: polish exportPreimage command

---
 cmd/geth/snapshot.go |  9 +--------
 cmd/utils/cmd.go     | 29 +++++++++++++++++++----------
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/cmd/geth/snapshot.go b/cmd/geth/snapshot.go
index da71dc94f3..481fbaa939 100644
--- a/cmd/geth/snapshot.go
+++ b/cmd/geth/snapshot.go
@@ -645,7 +645,6 @@ func snapshotExportPreimages(ctx *cli.Context) error {
 		}
 		root = headBlock.Root()
 	}
-
 	snapConfig := snapshot.Config{
 		CacheSize:  256,
 		Recovery:   false,
@@ -656,13 +655,7 @@ func snapshotExportPreimages(ctx *cli.Context) error {
 	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
+	return utils.ExportSnapshotPreimages(chaindb, snaptree, ctx.Args().First(), root)
 }
 
 // checkAccount iterates the snap data layers, and looks up the given account
diff --git a/cmd/utils/cmd.go b/cmd/utils/cmd.go
index 1af6ea5bea..4c265b74d8 100644
--- a/cmd/utils/cmd.go
+++ b/cmd/utils/cmd.go
@@ -386,14 +386,13 @@ func ExportSnapshotPreimages(chaindb ethdb.Database, snaptree *snapshot.Tree, fn
 	}
 	defer fh.Close()
 
+	// Enable gzip compressing if file name has gz suffix.
 	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
@@ -404,6 +403,11 @@ func ExportSnapshotPreimages(chaindb ethdb.Database, snaptree *snapshot.Tree, fn
 	}
 	hashCh := make(chan hashAndPreimageSize)
 
+	var (
+		start     = time.Now()
+		logged    = time.Now()
+		preimages int
+	)
 	go func() {
 		defer close(hashCh)
 		accIt, err := snaptree.AccountIterator(root, common.Hash{})
@@ -413,14 +417,14 @@ func ExportSnapshotPreimages(chaindb ethdb.Database, snaptree *snapshot.Tree, fn
 		}
 		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}
+			preimages += 1
+			hashCh <- hashAndPreimageSize{Hash: accIt.Hash(), Size: common.AddressLength}
 
 			if acc.Root != (common.Hash{}) && acc.Root != types.EmptyRootHash {
 				stIt, err := snaptree.StorageIterator(root, accIt.Hash(), common.Hash{})
@@ -429,13 +433,19 @@ func ExportSnapshotPreimages(chaindb ethdb.Database, snaptree *snapshot.Tree, fn
 					return
 				}
 				for stIt.Next() {
-					hashCh <- hashAndPreimageSize{Hash: stIt.Hash(), Size: 32}
+					preimages += 1
+					hashCh <- hashAndPreimageSize{Hash: stIt.Hash(), Size: common.HashLength}
+
+					if time.Since(logged) > time.Second*8 {
+						logged = time.Now()
+						log.Info("Exporting preimages", "count", preimages, "elapsed", common.PrettyDuration(time.Since(start)))
+					}
 				}
 				stIt.Release()
 			}
-			count++
-			if count%100000 == 0 {
-				log.Info("Last exported account", "account", accIt.Hash())
+			if time.Since(logged) > time.Second*8 {
+				logged = time.Now()
+				log.Info("Exporting preimages", "count", preimages, "elapsed", common.PrettyDuration(time.Since(start)))
 			}
 		}
 	}()
@@ -452,8 +462,7 @@ func ExportSnapshotPreimages(chaindb ethdb.Database, snaptree *snapshot.Tree, fn
 			return fmt.Errorf("failed to write preimage: %w", err)
 		}
 	}
-
-	log.Info("Exported preimages", "file", fn)
+	log.Info("Exported preimages", "count", preimages, "elapsed", common.PrettyDuration(time.Since(start)), "file", fn)
 	return nil
 }