From 26ca63f31b0d7d6ee3b948191764edf0a155788f Mon Sep 17 00:00:00 2001
From: Andre Medeiros <me@andremedeiros.info>
Date: Thu, 15 Jul 2021 20:09:38 -0400
Subject: [PATCH] add a profile option to the commands

---
 .gitignore           |  1 +
 internal/cmd/root.go | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+)

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))