From 2e39200214f7643c075369eb41b305ff7bcced3f Mon Sep 17 00:00:00 2001
From: c9s <yoanlin93@gmail.com>
Date: Fri, 15 Nov 2024 17:25:00 +0800
Subject: [PATCH] add pprof http server support with a build tag switch

---
 pkg/cmd/pprof.go |  6 ++++++
 pkg/cmd/root.go  | 23 +++++++++++++++++++++--
 2 files changed, 27 insertions(+), 2 deletions(-)
 create mode 100644 pkg/cmd/pprof.go

diff --git a/pkg/cmd/pprof.go b/pkg/cmd/pprof.go
new file mode 100644
index 000000000..e9223b2b4
--- /dev/null
+++ b/pkg/cmd/pprof.go
@@ -0,0 +1,6 @@
+//go:build pprof
+// +build pprof
+
+package cmd
+
+import _ "net/http/pprof"
diff --git a/pkg/cmd/root.go b/pkg/cmd/root.go
index 2b0ded1df..a696aef56 100644
--- a/pkg/cmd/root.go
+++ b/pkg/cmd/root.go
@@ -7,6 +7,7 @@ import (
 	"runtime/pprof"
 	"strings"
 	"time"
+	_ "time/tzdata"
 
 	"github.com/heroku/rollrus"
 	"github.com/joho/godotenv"
@@ -20,8 +21,6 @@ import (
 
 	"github.com/c9s/bbgo/pkg/bbgo"
 	"github.com/c9s/bbgo/pkg/util"
-
-	_ "time/tzdata"
 )
 
 var cpuProfileFile *os.File
@@ -81,6 +80,24 @@ var RootCmd = &cobra.Command{
 			}()
 		}
 
+		enableProfileServer, err := cmd.Flags().GetBool("enable-profile-server")
+		if err != nil {
+			return err
+		}
+
+		if enableProfileServer {
+			profileServerBind, err := cmd.Flags().GetString("profile-server-bind")
+			if err != nil {
+				return err
+			}
+
+			go func() {
+				if err := http.ListenAndServe(profileServerBind, nil); err != nil {
+					log.WithError(err).Errorf("profile server error")
+				}
+			}()
+		}
+
 		cpuProfile, err := cmd.Flags().GetString("cpu-profile")
 		if err != nil {
 			return err
@@ -195,6 +212,8 @@ func init() {
 	RootCmd.PersistentFlags().String("max-api-secret", "", "max api secret")
 
 	RootCmd.PersistentFlags().String("cpu-profile", "", "cpu profile")
+	RootCmd.PersistentFlags().Bool("enable-profile-server", false, "enable profile server binding")
+	RootCmd.PersistentFlags().String("profile-server-bind", "localhost:6060", "profile server binding")
 
 	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))