diff --git a/config.go b/config.go index ace9aec..7b41252 100644 --- a/config.go +++ b/config.go @@ -44,6 +44,7 @@ type Config struct { LocalAuth bool Password string AllowOrigins bool + PprofAddr string } func (cfg *Config) Parse(c *cli.Command) error { @@ -66,6 +67,7 @@ func (cfg *Config) Parse(c *cli.Command) error { getFlagOpt(c, "token", &cfg.Token) getFlagOpt(c, "password", &cfg.Password) getFlagOpt(c, "allow-origins", &cfg.AllowOrigins) + getFlagOpt(c, "pprof", &cfg.PprofAddr) return nil } diff --git a/main.go b/main.go index 5eb731c..b4f4c51 100644 --- a/main.go +++ b/main.go @@ -26,6 +26,7 @@ package main import ( "context" + _ "net/http/pprof" "os" "runtime" "runtime/debug" @@ -118,6 +119,10 @@ func main() { Name: "allow-origins", Usage: "allow all origins for cross-domain request", }, + &cli.StringFlag{ + Name: "pprof", + Usage: "enable pprof and listen on specified address (e.g. localhost:6060)", + }, &cli.BoolFlag{ Name: "verbose", Aliases: []string{"V"}, diff --git a/server.go b/server.go index c8a0c30..3fe14bb 100644 --- a/server.go +++ b/server.go @@ -25,6 +25,8 @@ package main import ( + "net" + "net/http" "sync" "sync/atomic" @@ -46,12 +48,39 @@ type DeviceGroup struct { func (srv *RttyServer) Run() error { log.Debug().Msgf("%+v", srv.cfg) + if srv.cfg.PprofAddr != "" { + go srv.ListenPprof() + } + go srv.ListenDevices() go srv.ListenHttpProxy() return srv.ListenAPI() } +func (srv *RttyServer) ListenPprof() { + ln, err := net.Listen("tcp", srv.cfg.PprofAddr) + if err != nil { + log.Error().Err(err).Msgf("Failed to start pprof server") + return + } + defer ln.Close() + + addr := ln.Addr().(*net.TCPAddr) + log.Info().Msgf("Starting pprof server on: %s", addr) + + host := addr.IP.String() + if host == "0.0.0.0" || host == "::" { + host = "localhost" + } + log.Info().Msgf("Access pprof at: http://%s:%d/debug/pprof/", host, addr.Port) + + err = http.Serve(ln, nil) + if err != nil { + log.Error().Err(err).Msgf("pprof server failed") + } +} + func (srv *RttyServer) GetDevice(group, id string) *Device { srv.mu.RLock() defer srv.mu.RUnlock()