Merge pull request #725 from senthilsam/feat/pprof

feat: add pprof support

Author: karimra

pkg/app/app.go

@@ -14,6 +14,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"net"
 	"net/http"
 	"os"
 	"sort"
@@ -111,6 +112,9 @@ type App struct {
 	tunTargetCfn  map[tunnel.Target]context.CancelFunc
 	// processors plugin manager
 	pm *plugin_manager.PluginManager
+
+	// pprof
+	pprof *pprofServer
 }
 
 func New() *App {
@@ -149,6 +153,9 @@ func New() *App {
 		ttm:          new(sync.RWMutex),
 		tunTargets:   make(map[tunnel.Target]struct{}),
 		tunTargetCfn: make(map[tunnel.Target]context.CancelFunc),
+
+		// pprof
+		pprof: newPprofServer(),
 	}
 	a.router.StrictSlash(true)
 	a.router.Use(headersMiddleware, a.loggingMiddleware)
@@ -177,6 +184,8 @@ func (a *App) InitGlobalFlags() {
 	a.RootCmd.PersistentFlags().StringVarP(&a.Config.GlobalFlags.TLSKey, "tls-key", "", "", "tls key")
 	a.RootCmd.PersistentFlags().DurationVarP(&a.Config.GlobalFlags.Timeout, "timeout", "", 10*time.Second, "grpc timeout, valid formats: 10s, 1m30s, 1h")
 	a.RootCmd.PersistentFlags().BoolVarP(&a.Config.GlobalFlags.Debug, "debug", "d", false, "debug mode")
+	a.RootCmd.PersistentFlags().BoolVarP(&a.Config.GlobalFlags.EnablePprof, "enable-pprof", "", false, "enable go pprof")
+	a.RootCmd.PersistentFlags().StringVarP(&a.Config.GlobalFlags.PprofAddr, "pprof-addr", "", defaultPprofAddr, "pprof host/IP and port to listen on")
 	a.RootCmd.PersistentFlags().BoolVarP(&a.Config.GlobalFlags.SkipVerify, "skip-verify", "", false, "skip verify tls connection")
 	a.RootCmd.PersistentFlags().BoolVarP(&a.Config.GlobalFlags.NoPrefix, "no-prefix", "", false, "do not add [ip:port] prefix to print output in case of multiple targets")
 	a.RootCmd.PersistentFlags().BoolVarP(&a.Config.GlobalFlags.ProxyFromEnv, "proxy-from-env", "", false, "use proxy from environment")
@@ -218,6 +227,19 @@ func (a *App) InitGlobalFlags() {
 }
 
 func (a *App) PreRunE(cmd *cobra.Command, args []string) error {
+	if a.Config.EnablePprof {
+		_, _, err := net.SplitHostPort(a.Config.GlobalFlags.PprofAddr)
+		if err != nil {
+			return fmt.Errorf("pprof error %v", err)
+		}
+		a.pprof.Start(a.Config.GlobalFlags.PprofAddr)
+		a.Logger.Printf("pprof server started at %s/debug/pprof", a.Config.GlobalFlags.PprofAddr)
+		go func() {
+			err := <-a.pprof.ErrChan()
+			a.Logger.Printf("pprof server failed: %v", err)
+		}()
+	}
+
 	a.Config.SetGlobalsFromEnv(a.RootCmd)
 	a.Config.SetPersistentFlagsFromFile(a.RootCmd)
 

pkg/app/const.go

@@ -14,6 +14,7 @@ const (
 	defaultGrpcPort   = "57400"
 	msgSize           = 512 * 1024 * 1024
 	defaultRetryTimer = 10 * time.Second
+	defaultPprofAddr  = "127.0.0.1:6060"
 
 	formatJSON      = "json"
 	formatPROTOJSON = "protojson"

pkg/app/pprof.go

@@ -0,0 +1,35 @@
+package app
+
+import (
+	"net/http"
+	_ "net/http/pprof" //nolint:gosec // Import for pprof, only enabled via CLI flag
+	"time"
+)
+
+type pprofServer struct {
+	err chan error
+}
+
+func newPprofServer() *pprofServer {
+	return &pprofServer{
+		err: make(chan error, 1),
+	}
+}
+
+func (p *pprofServer) Start(address string) {
+	go func() {
+		server := &http.Server{
+			Addr:              address,
+			ReadHeaderTimeout: 10 * time.Second,
+		}
+
+		if err := server.ListenAndServe(); err != nil {
+			p.err <- err
+		}
+		close(p.err)
+	}()
+}
+
+func (p *pprofServer) ErrChan() <-chan error {
+	return p.err
+}

pkg/config/config.go

@@ -91,6 +91,8 @@ type GlobalFlags struct {
 	LogTLSSecret  bool          `mapstructure:"log-tls-secret,omitempty" json:"log-tls-secret,omitempty" yaml:"log-tls-secret,omitempty"`
 	Timeout       time.Duration `mapstructure:"timeout,omitempty" json:"timeout,omitempty" yaml:"timeout,omitempty"`
 	Debug         bool          `mapstructure:"debug,omitempty" json:"debug,omitempty" yaml:"debug,omitempty"`
+	EnablePprof   bool          `mapstructure:"enable-pprof,omitempty" json:"enable-pprof,omitempty" yaml:"enable-pprof,omitempty"`
+	PprofAddr     string        `mapstructure:"pprof-addr,omitempty" json:"pprof-addr,omitempty" yaml:"pprof-addr,omitempty"`
 	SkipVerify    bool          `mapstructure:"skip-verify,omitempty" json:"skip-verify,omitempty" yaml:"skip-verify,omitempty"`
 	NoPrefix      bool          `mapstructure:"no-prefix,omitempty" json:"no-prefix,omitempty" yaml:"no-prefix,omitempty"`
 	ProxyFromEnv  bool          `mapstructure:"proxy-from-env,omitempty" json:"proxy-from-env,omitempty" yaml:"proxy-from-env,omitempty"`