From 0c2ac3a1684a9768e7d009e4f8711e1e80512160 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <me@cschwarz.com>
Date: Fri, 7 Sep 2018 00:04:03 -0700
Subject: [PATCH] pprof subcommand

---
 client/pprof.go | 35 +++++++++++++++++++++++++++++++++++
 main.go         | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)
 create mode 100644 client/pprof.go

diff --git a/client/pprof.go b/client/pprof.go
new file mode 100644
index 0000000..7afb4cb
--- /dev/null
+++ b/client/pprof.go
@@ -0,0 +1,35 @@
+package client
+
+import (
+	"github.com/zrepl/zrepl/config"
+	"github.com/zrepl/zrepl/daemon"
+	"log"
+	"os"
+)
+
+type PProfArgs struct {
+	daemon.PprofServerControlMsg
+}
+
+func RunPProf(conf *config.Config, args PProfArgs) {
+	log := log.New(os.Stderr, "", 0)
+
+	die := func() {
+		log.Printf("exiting after error")
+		os.Exit(1)
+	}
+
+	log.Printf("connecting to zrepl daemon")
+
+	httpc, err := controlHttpClient(conf.Global.Control.SockPath)
+	if err != nil {
+		log.Printf("error creating http client: %s", err)
+		die()
+	}
+	err = jsonRequestResponse(httpc, daemon.ControlJobEndpointPProf, args.PprofServerControlMsg, struct{}{})
+	if err != nil {
+		log.Printf("error sending control message: %s", err)
+		die()
+	}
+	log.Printf("finished")
+}
diff --git a/main.go b/main.go
index a41b080..9ab031d 100644
--- a/main.go
+++ b/main.go
@@ -113,6 +113,41 @@ var versionCmd = &cobra.Command{
 	},
 }
 
+var pprofCmd = &cobra.Command{
+	Use:   "pprof off | [on TCP_LISTEN_ADDRESS]",
+	Short: "start a http server exposing go-tool-compatible profiling endpoints at TCP_LISTEN_ADDRESS",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		conf, err := config.ParseConfig(rootArgs.configFile)
+		if err != nil {
+			return err
+		}
+
+		var pprofCmdArgs client.PProfArgs
+		if cmd.Flags().NArg() < 1 {
+			goto enargs
+		}
+		switch cmd.Flags().Arg(0) {
+		case "on":
+			pprofCmdArgs.Run = true
+			if cmd.Flags().NArg() != 2 {
+				return errors.New("must specify TCP_LISTEN_ADDRESS as second positional argument")
+			}
+			pprofCmdArgs.HttpListenAddress = cmd.Flags().Arg(1)
+		case "off":
+			if cmd.Flags().NArg() != 1 {
+				goto enargs
+			}
+			pprofCmdArgs.Run = false
+		}
+
+		client.RunPProf(conf, pprofCmdArgs)
+		return nil
+	enargs:
+		return errors.New("invalid number of positional arguments")
+
+	},
+}
+
 var rootArgs struct {
 	configFile string
 }
@@ -128,6 +163,7 @@ func init() {
 	rootCmd.AddCommand(configcheckCmd)
 	versionCmd.Flags().StringVar(&versionCmdArgs.Show, "show", "", "version info to show (client|daemon)")
 	rootCmd.AddCommand(versionCmd)
+	rootCmd.AddCommand(pprofCmd)
 }
 
 func main() {