From f8dbd267368112b0a44259ac61b9f4202f18e38f Mon Sep 17 00:00:00 2001
From: c9s <yoanlin93@gmail.com>
Date: Thu, 9 Jun 2022 15:49:52 +0800
Subject: [PATCH] move cpu profile option to global cmd

---
 pkg/cmd/backtest.go |  1 -
 pkg/cmd/root.go     | 32 +++++++++++++++++++++++++++++++-
 pkg/cmd/run.go      |  1 -
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/pkg/cmd/backtest.go b/pkg/cmd/backtest.go
index 7771d2c6a..92205c3ce 100644
--- a/pkg/cmd/backtest.go
+++ b/pkg/cmd/backtest.go
@@ -127,7 +127,6 @@ var BacktestCmd = &cobra.Command{
 			return err
 		}
 
-
 		if userConfig.Backtest == nil {
 			return errors.New("backtest config is not defined")
 		}
diff --git a/pkg/cmd/root.go b/pkg/cmd/root.go
index 694440e71..0bfec7e66 100644
--- a/pkg/cmd/root.go
+++ b/pkg/cmd/root.go
@@ -4,6 +4,7 @@ import (
 	"net/http"
 	"os"
 	"path"
+	"runtime/pprof"
 	"strings"
 	"time"
 
@@ -22,6 +23,8 @@ import (
 	_ "github.com/go-sql-driver/mysql"
 )
 
+var cpuProfileFile *os.File
+
 var userConfig *bbgo.Config
 
 var RootCmd = &cobra.Command{
@@ -32,7 +35,7 @@ var RootCmd = &cobra.Command{
 	SilenceUsage: true,
 
 	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
-		if err := cobraLoadDotenv(cmd, args) ; err != nil {
+		if err := cobraLoadDotenv(cmd, args); err != nil {
 			return err
 		}
 
@@ -53,8 +56,34 @@ var RootCmd = &cobra.Command{
 			}()
 		}
 
+		cpuProfile, err := cmd.Flags().GetString("cpu-profile")
+		if err != nil {
+			return err
+		}
+
+		if cpuProfile != "" {
+			log.Infof("starting cpu profiler...")
+
+			cpuProfileFile, err = os.Create(cpuProfile)
+			if err != nil {
+				log.Fatal("could not create CPU profile: ", err)
+			}
+
+			if err := pprof.StartCPUProfile(cpuProfileFile); err != nil {
+				log.Fatal("could not start CPU profile: ", err)
+			}
+		}
+
 		return cobraLoadConfig(cmd, args)
 	},
+	PersistentPostRunE: func(cmd *cobra.Command, args []string) error {
+		pprof.StopCPUProfile()
+		if cpuProfileFile != nil {
+			return cpuProfileFile.Close() // error handling omitted for example
+		}
+
+		return nil
+	},
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return nil
 	},
@@ -139,6 +168,7 @@ func init() {
 	RootCmd.PersistentFlags().String("ftx-api-key", "", "ftx api key")
 	RootCmd.PersistentFlags().String("ftx-api-secret", "", "ftx api secret")
 	RootCmd.PersistentFlags().String("ftx-subaccount", "", "subaccount name. Specify it if the credential is for subaccount.")
+	RootCmd.PersistentFlags().String("cpu-profile", "", "cpu profile")
 
 	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
 
diff --git a/pkg/cmd/run.go b/pkg/cmd/run.go
index 9dd4c4014..4d769b2bb 100644
--- a/pkg/cmd/run.go
+++ b/pkg/cmd/run.go
@@ -34,7 +34,6 @@ func init() {
 	RunCmd.Flags().Bool("enable-grpc", false, "enable grpc server")
 	RunCmd.Flags().String("grpc-bind", ":50051", "grpc server binding")
 
-	RunCmd.Flags().String("cpu-profile", "", "cpu profile")
 	RunCmd.Flags().Bool("setup", false, "use setup mode")
 	RootCmd.AddCommand(RunCmd)
 }