diff --git a/.gitignore b/.gitignore index 1c24a10..e42a7de 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ tmp/ .env dump.rdb +cpu.pprof diff --git a/internal/cmd/root.go b/internal/cmd/root.go index 8204d41..b2843c9 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -2,6 +2,8 @@ package cmd import ( "context" + "os" + "runtime/pprof" "github.com/DataDog/datadog-go/statsd" "github.com/adjust/rmq/v4" @@ -25,11 +27,31 @@ type Command struct { func Execute(ctx context.Context) int { _ = godotenv.Load() + profile := false + rootCmd := &cobra.Command{ Use: "apollo", Short: "Apollo is an amazing Reddit client for iOS. This isn't it, but it helps.", + PersistentPreRunE: func(cmd *cobra.Command, args []string) error { + if !profile { + return nil + } + f, perr := os.Create("cpu.pprof") + if perr != nil { + return perr + } + pprof.StartCPUProfile(f) + return nil + }, + PersistentPostRun: func(cmd *cobra.Command, args []string) { + if profile { + pprof.StopCPUProfile() + } + }, } + rootCmd.PersistentFlags().BoolVarP(&profile, "profile", "p", false, "record CPU pprof") + rootCmd.AddCommand(APICmd(ctx)) rootCmd.AddCommand(SchedulerCmd(ctx)) rootCmd.AddCommand(WorkerCmd(ctx))