diff --git a/cmd/violet/serve.go b/cmd/violet/serve.go index 1aec31f..a97a78f 100644 --- a/cmd/violet/serve.go +++ b/cmd/violet/serve.go @@ -25,17 +25,22 @@ import ( _ "net/http/pprof" "os" "path/filepath" + "runtime/pprof" ) -type serveCmd struct{ configPath string } +type serveCmd struct { + configPath string + cpuprofile string +} func (s *serveCmd) Name() string { return "serve" } func (s *serveCmd) Synopsis() string { return "Serve reverse proxy server" } func (s *serveCmd) SetFlags(f *flag.FlagSet) { f.StringVar(&s.configPath, "conf", "", "/path/to/config.json : path to the config file") + f.StringVar(&s.cpuprofile, "cpuprofile", "", "write cpu profile to file") } func (s *serveCmd) Usage() string { - return `serve [-conf ] + return `serve [-conf ] [-cpuprofile ] Serve reverse proxy server using information from config file ` } @@ -43,6 +48,17 @@ func (s *serveCmd) Usage() string { func (s *serveCmd) Execute(_ context.Context, _ *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus { log.Println("[Violet] Starting...") + // Enable cpu profiling + if s.cpuprofile != "" { + f, err := os.Create(s.cpuprofile) + if err != nil { + log.Fatal(err) + } + log.Printf("[Violet] CPU profiling enabled, writing to '%s'\n", s.cpuprofile) + _ = pprof.StartCPUProfile(f) + defer pprof.StopCPUProfile() + } + if s.configPath == "" { log.Println("[Violet] Error: config flag is missing") return subcommands.ExitUsageError