From 929348014592abb28d18d274cb12729e766f6e41 Mon Sep 17 00:00:00 2001
From: cwbhhjl <cwbhhjl@gmail.com>
Date: Tue, 11 Jan 2022 07:52:51 +0800
Subject: [PATCH] =?UTF-8?q?Merge=20PR=EF=BC=9Asupport=20gen=20mem=20pprof?=
 =?UTF-8?q?=20(#1405)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: KamiD <44460798+KamiD@users.noreply.github.com>
---
 cmd/exchaind/replay.go | 32 ++++++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/cmd/exchaind/replay.go b/cmd/exchaind/replay.go
index ec2c2d0a3a..745e7e8dc3 100644
--- a/cmd/exchaind/replay.go
+++ b/cmd/exchaind/replay.go
@@ -7,6 +7,7 @@ import (
 	_ "net/http/pprof"
 	"os"
 	"path/filepath"
+	"runtime"
 	"runtime/pprof"
 	"time"
 
@@ -38,8 +39,9 @@ const (
 	blockStoreDB     = "blockstore"
 	stateDB          = "state"
 
-	pprofAddrFlag    = "pprof_addr"
-	runWithPprofFlag = "gen_pprof"
+	pprofAddrFlag       = "pprof_addr"
+	runWithPprofFlag    = "gen_pprof"
+	runWithPprofMemFlag = "gen_pprof_mem"
 
 	saveBlock = "save_block"
 
@@ -65,6 +67,17 @@ func replayCmd(ctx *server.Context) *cobra.Command {
 			replayBlock(ctx, dataDir)
 			log.Println("--------- replay success ---------")
 		},
+		PostRun: func(cmd *cobra.Command, args []string) {
+			if viper.GetBool(runWithPprofMemFlag) {
+				log.Println("--------- gen pprof mem start ---------")
+				err := dumpMemPprof()
+				if err != nil {
+					log.Println(err)
+				} else {
+					log.Println("--------- gen pprof mem success ---------")
+				}
+			}
+		},
 	}
 	cmd.Flags().StringP(replayedBlockDir, "d", ".exchaind/data", "Directory of block data to be replayed")
 	cmd.Flags().StringP(pprofAddrFlag, "p", "0.0.0.0:26661", "Address and port of pprof HTTP server listening")
@@ -97,6 +110,7 @@ func replayCmd(ctx *server.Context) *cobra.Command {
 	cmd.Flags().BoolVar(&tmiavl.EnableAsyncCommit, tmiavl.FlagIavlEnableAsyncCommit, false, "Enable cache iavl node data to optimization leveldb pruning process")
 	cmd.Flags().BoolVar(&system.EnableGid, system.FlagEnableGid, false, "Display goroutine id in log")
 	cmd.Flags().Bool(runWithPprofFlag, false, "Dump the pprof of the entire replay process")
+	cmd.Flags().Bool(runWithPprofMemFlag, false, "Dump the mem profile of the entire replay process")
 	cmd.Flags().Bool(sm.FlagParalleledTx, false, "pall Tx")
 	cmd.Flags().Bool(saveBlock, false, "save block when replay")
 	cmd.Flags().Int64(config.FlagMaxGasUsedPerBlock, -1, "Maximum gas used of transactions in a block")
@@ -292,6 +306,20 @@ func doReplay(ctx *server.Context, state sm.State, stateStoreDB dbm.DB,
 	}
 }
 
+func dumpMemPprof() error {
+	fileName := fmt.Sprintf("replay_pprof_%s.mem.bin", time.Now().Format("20060102150405"))
+	f, err := os.Create(fileName)
+	if err != nil {
+		return fmt.Errorf("create mem pprof file %s error: %w", fileName, err)
+	}
+	defer f.Close()
+	runtime.GC() // get up-to-date statistics
+	if err = pprof.WriteHeapProfile(f); err != nil {
+		return fmt.Errorf("could not write memory profile: %w", err)
+	}
+	return nil
+}
+
 func startDumpPprof() {
 	var (
 		binarySuffix = time.Now().Format("20060102150405") + ".bin"