From 9d207b5e9a5a90a922b356ebc0b9ee647752edf8 Mon Sep 17 00:00:00 2001
From: gray <greyschwinger@gmail.com>
Date: Mon, 10 Jun 2024 02:34:19 +0800
Subject: [PATCH] Add global config to enable pprof

---
 cmd/run.go       | 19 +++++++++++++++++++
 config/config.go |  1 +
 example.dae      |  3 +++
 3 files changed, 23 insertions(+)

diff --git a/cmd/run.go b/cmd/run.go
index 633593f15..019539324 100644
--- a/cmd/run.go
+++ b/cmd/run.go
@@ -25,6 +25,8 @@ import (
 	"github.com/daeuniverse/outbound/protocol/direct"
 	"gopkg.in/natefinch/lumberjack.v2"
 
+	_ "net/http/pprof"
+
 	"github.com/daeuniverse/dae/cmd/internal"
 	"github.com/daeuniverse/dae/common"
 	"github.com/daeuniverse/dae/common/consts"
@@ -124,6 +126,13 @@ func Run(log *logrus.Logger, conf *config.Config, externGeoDataDirs []string) (e
 		return err
 	}
 
+	var pprofServer *http.Server
+	if conf.Global.PprofPort != 0 {
+		pprofAddr := fmt.Sprintf("localhost:%d", conf.Global.PprofPort)
+		pprofServer = &http.Server{Addr: pprofAddr, Handler: nil}
+		go pprofServer.ListenAndServe()
+	}
+
 	// Serve tproxy TCP/UDP server util signals.
 	var listener *control.Listener
 	sigs := make(chan os.Signal, 1)
@@ -270,6 +279,16 @@ loop:
 				oldC.AbortConnections()
 			}
 			oldC.Close()
+
+			if pprofServer != nil {
+				pprofServer.Shutdown(context.Background())
+				pprofServer = nil
+			}
+			if newConf.Global.PprofPort != 0 {
+				pprofAddr := fmt.Sprintf("localhost:%d", conf.Global.PprofPort)
+				pprofServer = &http.Server{Addr: pprofAddr, Handler: nil}
+				go pprofServer.ListenAndServe()
+			}
 		case syscall.SIGHUP:
 			// Ignore.
 			continue
diff --git a/config/config.go b/config/config.go
index e9cb944ca..4031e4ee9 100644
--- a/config/config.go
+++ b/config/config.go
@@ -41,6 +41,7 @@ type Global struct {
 	SniffingTimeout        time.Duration `mapstructure:"sniffing_timeout" default:"100ms"`
 	TlsImplementation      string        `mapstructure:"tls_implementation" default:"tls"`
 	UtlsImitate            string        `mapstructure:"utls_imitate" default:"chrome_auto"`
+	PprofPort              uint16        `mapstructure:"pprof_port" default:"0"`
 }
 
 type Utls struct {
diff --git a/example.dae b/example.dae
index 6bf5746f1..1a4740ced 100644
--- a/example.dae
+++ b/example.dae
@@ -93,6 +93,9 @@ global {
     # The Client Hello ID for uTLS to imitate. This takes effect only if tls_implementation is utls.
     # See more: https://github.com/daeuniverse/dae/blob/331fa23c16/component/outbound/transport/tls/utls.go#L17
     utls_imitate: chrome_auto
+
+    # Set non-zero value to enable pprof.
+    pprof_port: 6060
 }
 
 # Subscriptions defined here will be resolved as nodes and merged as a part of the global node pool.