diff --git a/main.go b/main.go index 51afe11..7d95abf 100644 --- a/main.go +++ b/main.go @@ -1,19 +1,20 @@ package main import ( - "fmt" - "bufio" "encoding/json" + "fmt" "net/http" "os" "strconv" + "github.com/prometheus/client_golang/prometheus/promhttp" + "github.com/spf13/cobra" + + "github.com/arangodb/feed/pkg/cli" "github.com/arangodb/feed/pkg/config" "github.com/arangodb/feed/pkg/feedlang" "github.com/arangodb/feed/pkg/operations" - "github.com/prometheus/client_golang/prometheus/promhttp" - "github.com/spf13/cobra" ) const ( @@ -116,5 +117,7 @@ func mainExecute(cmd *cobra.Command, _ []string) error { } func main() { - cmd.Execute() + if err := cli.RunCommandWithProfile(cmd); err != nil { + os.Exit(1) + } } diff --git a/pkg/cli/profiling.go b/pkg/cli/profiling.go new file mode 100644 index 0000000..98b27c8 --- /dev/null +++ b/pkg/cli/profiling.go @@ -0,0 +1,64 @@ +package cli + +import ( + "fmt" + "os" + "runtime" + "runtime/pprof" + + "github.com/pkg/errors" + "github.com/spf13/cobra" +) + +// RunCommandWithProfile enables profile for CPU/Memory if env variables requests so +func RunCommandWithProfile(cmd *cobra.Command) error { + // Turn on CPU profiling before the command is launched. + if filename := os.Getenv("PPROF_CPU_FILENAME"); len(filename) > 0 { + f, err := os.Create(filename) + if err != nil { + fmt.Printf("could not create CPU profile filename %s\n", filename) + return nil + } + defer f.Close() + + if err := pprof.StartCPUProfile(f); err != nil { + fmt.Println("could not start CPU profile") + return nil + } + defer pprof.StopCPUProfile() + } + + errExecute := cmd.Execute() + + // Close memory profiling after the command. + if err := profileMemory(os.Getenv("PPROF_MEMORY_FILENAME")); err != nil { + fmt.Printf("profileMemory failed: %s\n", err) + if errExecute == nil { + // When command does not fail then memory profile error can be returned. + return err + } + } + + return errExecute +} + +// profileMemory reports memory usage in the given file. +func profileMemory(filename string) error { + if len(filename) == 0 { + return nil + } + + f, err := os.Create(filename) + if err != nil { + return errors.WithMessagef(err, "could not create memory profile filename %s", filename) + } + defer f.Close() + + // Get up-to-date statistics. + runtime.GC() + if err := pprof.WriteHeapProfile(f); err != nil { + return errors.WithMessagef(err, "could not write memory profile for filename %s", filename) + } + + return nil +}