diff --git a/cmd/pilosa/main.go b/cmd/pilosa/main.go index fd26a5685..7107745fe 100644 --- a/cmd/pilosa/main.go +++ b/cmd/pilosa/main.go @@ -9,6 +9,7 @@ import ( "os" "os/signal" "path/filepath" + "runtime/pprof" "strings" "time" @@ -45,18 +46,38 @@ func main() { os.Exit(2) } + // Start CPU profiling. + if m.CPUProfile != "" { + f, err := os.Create(m.CPUProfile) + if err != nil { + fmt.Fprintf(m.Stderr, "create cpu profile: %v", err) + os.Exit(1) + } + defer f.Close() + + fmt.Fprintln(m.Stderr, "Starting cpu profile") + pprof.StartCPUProfile(f) + time.AfterFunc(m.CPUTime, func() { + fmt.Fprintln(m.Stderr, "Stopping cpu profile") + pprof.StopCPUProfile() + f.Close() + }) + } + // Execute the program. if err := m.Run(); err != nil { fmt.Fprintln(m.Stderr, err) + fmt.Fprintln(m.Stderr, "stopping profile") os.Exit(1) } // First SIGKILL causes server to shut down gracefully. - // Second signal causes a hard shutdown. c := make(chan os.Signal, 2) signal.Notify(c, os.Interrupt) sig := <-c fmt.Fprintf(m.Stderr, "Received %s; gracefully shutting down...\n", sig.String()) + + // Second signal causes a hard shutdown. go func() { <-c; os.Exit(1) }() if err := m.Close(); err != nil { @@ -73,6 +94,10 @@ type Main struct { ConfigPath string Config *Config + // Profiling options. + CPUProfile string + CPUTime time.Duration + // Standard input/output Stdin io.Reader Stdout io.Writer @@ -127,8 +152,10 @@ func (m *Main) Close() error { // ParseFlags parses command line flags from args. func (m *Main) ParseFlags(args []string) error { fs := flag.NewFlagSet("pilosa", flag.ContinueOnError) - fs.SetOutput(m.Stderr) + fs.StringVar(&m.CPUProfile, "cpuprofile", "", "cpu profile") + fs.DurationVar(&m.CPUTime, "cputime", 30*time.Second, "cpu profile duration") fs.StringVar(&m.ConfigPath, "config", "", "config path") + fs.SetOutput(m.Stderr) if err := fs.Parse(args); err != nil { return err }