diff --git a/cmd/server/main.go b/cmd/server/main.go index 1edab552e..6fe79593c 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -10,6 +10,8 @@ import ( "net/http" "os" "os/signal" + "runtime" + "runtime/pprof" "sync" "syscall" "time" @@ -46,6 +48,14 @@ func main() { Usage: "API keys/secret pairs (key:secret, one per line)", EnvVars: []string{"KEY_FILE"}, }, + &cli.StringFlag{ + Name: "cpuprofile", + Usage: "write cpu profile to `file`", + }, + &cli.StringFlag{ + Name: "memprofile", + Usage: "write memory profile to `file`", + }, &cli.BoolFlag{ Name: "dev", Usage: "when set, token validation will be disabled", @@ -69,6 +79,8 @@ func main() { func startServer(c *cli.Context) error { rand.Seed(time.Now().UnixNano()) + cpuProfile := c.String("cpuprofile") + memProfile := c.String("memprofile") conf, err := config.NewConfig(c.String("config")) if err != nil { return err @@ -83,6 +95,31 @@ func startServer(c *cli.Context) error { logger.InitProduction() } + if cpuProfile != "" { + if f, err := os.Create(cpuProfile); err != nil { + return err + } else { + defer f.Close() + if err := pprof.StartCPUProfile(f); err != nil { + return err + } + defer pprof.StopCPUProfile() + } + } + + if memProfile != "" { + if f, err := os.Create(memProfile); err != nil { + return err + } else { + defer func() { + // run memory profile at termination + runtime.GC() + pprof.WriteHeapProfile(f) + f.Close() + }() + } + } + // require a key provider if keyProvider, err = createKeyProvider(c.String("key-file"), c.String("keys")); err != nil { return err