pprof subcommand

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")
+}

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