feat(server/v2): enable pprof (#22284)

server/v2/commands.go

@@ -6,13 +6,15 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"runtime/pprof"
 	"strings"
 	"syscall"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 
 	"cosmossdk.io/core/transaction"
+	"cosmossdk.io/log"
 )
 
 // Execute executes the root command of an application.
@@ -127,11 +129,9 @@ func createStartCommand[T transaction.Tx](
 				}
 			}()
 
-			if err := server.Start(ctx); err != nil {
-				return err
-			}
-
-			return nil
+			return wrapCPUProfile(l, v, func() error {
+				return server.Start(ctx)
+			})
 		},
 	}
 
@@ -143,6 +143,37 @@ func createStartCommand[T transaction.Tx](
 	return cmd
 }
 
+// wrapCPUProfile starts CPU profiling, if enabled, and executes the provided
+// callbackFn, then waits for it to return.
+func wrapCPUProfile(logger log.Logger, v *viper.Viper, callbackFn func() error) error {
+	cpuProfileFile := v.GetString(FlagCPUProfiling)
+	if len(cpuProfileFile) == 0 {
+		// if cpu profiling is not enabled, just run the callback
+		return callbackFn()
+	}
+
+	f, err := os.Create(cpuProfileFile)
+	if err != nil {
+		return err
+	}
+
+	logger.Info("starting CPU profiler", "profile", cpuProfileFile)
+	if err := pprof.StartCPUProfile(f); err != nil {
+		_ = f.Close()
+		return err
+	}
+
+	defer func() {
+		logger.Info("stopping CPU profiler", "profile", cpuProfileFile)
+		pprof.StopCPUProfile()
+		if err := f.Close(); err != nil {
+			logger.Info("failed to close cpu-profile file", "profile", cpuProfileFile, "err", err.Error())
+		}
+	}()
+
+	return callbackFn()
+}
+
 // configHandle writes the default config to the home directory if it does not exist and sets the server context
 func configHandle[T transaction.Tx](s *Server[T], cmd *cobra.Command) error {
 	home, err := cmd.Flags().GetString(FlagHome)

server/v2/flags.go

@@ -9,7 +9,10 @@ func prefix(f string) string {
 	return fmt.Sprintf("%s.%s", serverName, f)
 }
 
-var FlagMinGasPrices = prefix("minimum-gas-prices")
+var (
+	FlagMinGasPrices = prefix("minimum-gas-prices")
+	FlagCPUProfiling = prefix("cpu-profile")
+)
 
 const (
 	// FlagHome specifies the home directory flag.

server/v2/server.go

@@ -181,6 +181,8 @@ func (s *Server[T]) Configs() map[string]any {
 func (s *Server[T]) StartCmdFlags() *pflag.FlagSet {
 	flags := pflag.NewFlagSet(s.Name(), pflag.ExitOnError)
 	flags.String(FlagMinGasPrices, "", "Minimum gas prices to accept for transactions; Any fee in a tx must meet this minimum (e.g. 0.01photino;0.0001stake)")
+	flags.String(FlagCPUProfiling, "", "Enable CPU profiling and write to the specified file")
+
 	return flags
 }