feat: support profiling block replay during abci handshake

CHANGELOG.md

@@ -144,6 +144,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#14406](https://github.com/cosmos/cosmos-sdk/issues/14406) Migrate usage of types/store.go to store/types/..
 * (x/staking) [#14590](https://github.com/cosmos/cosmos-sdk/pull/14590) Return undelegate amount in MsgUndelegateResponse
 * (tools) [#14793](https://github.com/cosmos/cosmos-sdk/pull/14793) Dockerfile optimization.
+* (cli) [#14953](https://github.com/cosmos/cosmos-sdk/pull/14953) Enable profiling block replay during abci handshake with `--cpu-profile`.
 
 ### State Machine Breaking
 

server/start.go

@@ -134,11 +134,15 @@ is performed. Note, when enabled, gRPC will also be automatically enabled.
 			withTM, _ := cmd.Flags().GetBool(flagWithTendermint)
 			if !withTM {
 				serverCtx.Logger.Info("starting ABCI without CometBFT")
-				return startStandAlone(serverCtx, appCreator)
+				return wrapCPUProfile(serverCtx, func() error {
+					return startStandAlone(serverCtx, appCreator)
+				})
 			}
 
 			// amino is needed here for backwards compatibility of REST routes
-			err = startInProcess(serverCtx, clientCtx, appCreator)
+			err = wrapCPUProfile(serverCtx, func() error {
+				return startInProcess(serverCtx, clientCtx, appCreator)
+			})
 			errCode, ok := err.(ErrorCode)
 			if !ok {
 				return err
@@ -249,27 +253,6 @@ func startStandAlone(ctx *Context, appCreator types.AppCreator) error {
 func startInProcess(ctx *Context, clientCtx client.Context, appCreator types.AppCreator) error {
 	cfg := ctx.Config
 	home := cfg.RootDir
-	var cpuProfileCleanup func()
-
-	if cpuProfile := ctx.Viper.GetString(flagCPUProfile); cpuProfile != "" {
-		f, err := os.Create(cpuProfile)
-		if err != nil {
-			return err
-		}
-
-		ctx.Logger.Info("starting CPU profiler", "profile", cpuProfile)
-		if err := pprof.StartCPUProfile(f); err != nil {
-			return err
-		}
-
-		cpuProfileCleanup = func() {
-			ctx.Logger.Info("stopping CPU profiler", "profile", cpuProfile)
-			pprof.StopCPUProfile()
-			if err := f.Close(); err != nil {
-				ctx.Logger.Info("failed to close cpu-profile file", "profile", cpuProfile, "err", err.Error())
-			}
-		}
-	}
 
 	db, err := openDB(home, GetAppDBBackend(ctx.Viper))
 	if err != nil {
@@ -282,16 +265,11 @@ func startInProcess(ctx *Context, clientCtx client.Context, appCreator types.App
 		return err
 	}
 
-	// Clean up the traceWriter in the cpuProfileCleanup routine that is invoked
-	// when the server is shutting down.
-	fn := cpuProfileCleanup
-	cpuProfileCleanup = func() {
-		if fn != nil {
-			fn()
-		}
-
-		// if flagTraceStore is not used then traceWriter is nil
-		if traceWriter != nil {
+	// Clean up the traceWriter when the server is shutting down.
+	var traceWriterCleanup func()
+	// if flagTraceStore is not used then traceWriter is nil
+	if traceWriter != nil {
+		traceWriterCleanup = func() {
 			if err = traceWriter.Close(); err != nil {
 				ctx.Logger.Error("failed to close trace writer", "err", err)
 			}
@@ -464,8 +442,8 @@ func startInProcess(ctx *Context, clientCtx client.Context, appCreator types.App
 			_ = tmNode.Stop()
 		}
 
-		if cpuProfileCleanup != nil {
-			cpuProfileCleanup()
+		if traceWriterCleanup != nil {
+			traceWriterCleanup()
 		}
 
 		if apiSrv != nil {
@@ -485,3 +463,40 @@ func startTelemetry(cfg serverconfig.Config) (*telemetry.Metrics, error) {
 	}
 	return telemetry.New(cfg.Telemetry)
 }
+
+// wrapCPUProfile runs callback in a goroutine, then wait for quit signals.
+func wrapCPUProfile(ctx *Context, callback func() error) error {
+	if cpuProfile := ctx.Viper.GetString(flagCPUProfile); cpuProfile != "" {
+		f, err := os.Create(cpuProfile)
+		if err != nil {
+			return err
+		}
+
+		ctx.Logger.Info("starting CPU profiler", "profile", cpuProfile)
+		if err := pprof.StartCPUProfile(f); err != nil {
+			return err
+		}
+
+		defer func() {
+			ctx.Logger.Info("stopping CPU profiler", "profile", cpuProfile)
+			pprof.StopCPUProfile()
+			if err := f.Close(); err != nil {
+				ctx.Logger.Info("failed to close cpu-profile file", "profile", cpuProfile, "err", err.Error())
+			}
+		}()
+	}
+
+	errCh := make(chan error)
+	go func() {
+		errCh <- callback()
+	}()
+
+	select {
+	case err := <-errCh:
+		return err
+
+	case <-time.After(types.ServerStartTime):
+	}
+
+	return WaitForQuitSignals()
+}