control pprof rewrite: expose pprof metrics via HTTP server controlled from CLI

cmd/config_job_control.go

@@ -8,7 +8,6 @@ import (
 	"io"
 	"net"
 	"net/http"
-	"net/http/pprof"
 )
 
 type ControlJob struct {
@@ -37,7 +36,7 @@ func (j *ControlJob) JobStatus(ctx context.Context) (*JobStatus, error) {
 }
 
 const (
-	ControlJobEndpointProfile string = "/debug/pprof/profile"
+	ControlJobEndpointPProf   string = "/debug/pprof"
 	ControlJobEndpointVersion string = "/version"
 	ControlJobEndpointStatus  string = "/status"
 )
@@ -55,8 +54,19 @@ func (j *ControlJob) JobStart(ctx context.Context) {
 		return
 	}
 
+	pprofServer := NewPProfServer(ctx)
+
 	mux := http.NewServeMux()
-	mux.Handle(ControlJobEndpointProfile, requestLogger{log: log, handlerFunc: pprof.Profile})
+	mux.Handle(ControlJobEndpointPProf, requestLogger{log: log, handlerFunc: func(w http.ResponseWriter, r *http.Request) {
+		var msg PprofServerControlMsg
+		err := json.NewDecoder(r.Body).Decode(&msg)
+		if err != nil {
+			log.WithError(err).Error("bad pprof request from client")
+			w.WriteHeader(http.StatusBadRequest)
+		}
+		pprofServer.Control(msg)
+		w.WriteHeader(200)
+	}})
 	mux.Handle(ControlJobEndpointVersion,
 		requestLogger{log: log, handler: jsonResponder{func() (interface{}, error) {
 			return NewZreplVersionInformation(), nil

cmd/control.go

@@ -6,13 +6,13 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/dustin/go-humanize"
+	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"github.com/zrepl/zrepl/logger"
 	"io"
 	golog "log"
 	"net"
 	"net/http"
-	"net/url"
 	"os"
 	"sort"
 	"strings"
@@ -25,12 +25,34 @@ var controlCmd = &cobra.Command{
 }
 
 var pprofCmd = &cobra.Command{
-	Use:   "pprof cpu OUTFILE",
-	Short: "pprof CPU of daemon to OUTFILE (- for stdout)",
+	Use:   "pprof off | [on TCP_LISTEN_ADDRESS]",
+	Short: "start a http server exposing go-tool-compatible profiling endpoints at TCP_LISTEN_ADDRESS",
 	Run:   doControlPProf,
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		if cmd.Flags().NArg() < 1 {
+			goto enargs
+		}
+		switch cmd.Flags().Arg(0) {
+		case "on":
+			pprofCmdArgs.msg.Run = true
+			if cmd.Flags().NArg() != 2 {
+				return errors.New("must specify TCP_LISTEN_ADDRESS as second positional argument")
+			}
+			pprofCmdArgs.msg.HttpListenAddress = cmd.Flags().Arg(1)
+		case "off":
+			if cmd.Flags().NArg() != 1 {
+				goto enargs
+			}
+			pprofCmdArgs.msg.Run = false
+		}
+		return nil
+	enargs:
+		return errors.New("invalid number of positional arguments")
+
+	},
 }
 var pprofCmdArgs struct {
-	seconds int64
+	msg PprofServerControlMsg
 }
 
 var controlVersionCmd = &cobra.Command{
@@ -54,7 +76,6 @@ var controlStatusCmd = &cobra.Command{
 func init() {
 	RootCmd.AddCommand(controlCmd)
 	controlCmd.AddCommand(pprofCmd)
-	pprofCmd.Flags().Int64Var(&pprofCmdArgs.seconds, "seconds", 30, "seconds to profile")
 	controlCmd.AddCommand(controlVersionCmd)
 	controlCmd.AddCommand(controlStatusCmd)
 	controlStatusCmd.Flags().StringVar(&controlStatusCmdArgs.format, "format", "human", "output format (human|raw)")
@@ -87,55 +108,25 @@ func doControlPProf(cmd *cobra.Command, args []string) {
 		os.Exit(1)
 	}
 
-	if cmd.Flags().Arg(0) != "cpu" {
-		log.Printf("only CPU profiles are supported")
-		log.Printf("%s", cmd.UsageString())
-		die()
-	}
-
-	outfn := cmd.Flags().Arg(1)
-	if outfn == "" {
-		log.Printf("must specify output filename")
-		log.Printf("%s", cmd.UsageString())
-		die()
-	}
-	var out io.Writer
-	var err error
-	if outfn == "-" {
-		out = os.Stdout
-	} else {
-		out, err = os.Create(outfn)
-		if err != nil {
-			log.Printf("error creating output file: %s", err)
-			die()
-		}
-	}
-
-	log.Printf("connecting to daemon")
+	log.Printf("connecting to zrepl daemon")
 	httpc, err := controlHttpClient()
 	if err != nil {
 		log.Printf("error parsing config: %s", err)
 		die()
 	}
 
-	log.Printf("profiling...")
-	v := url.Values{}
-	v.Set("seconds", fmt.Sprintf("%d", pprofCmdArgs.seconds))
-	v.Encode()
-	resp, err := httpc.Get("http://unix" + ControlJobEndpointProfile + "?" + v.Encode())
+	var buf bytes.Buffer
+	if err := json.NewEncoder(&buf).Encode(&pprofCmdArgs.msg); err != nil {
+		log.Printf("error marshaling request: %s", err)
+		die()
+	}
+	_, err = httpc.Post("http://unix"+ControlJobEndpointPProf, "application/json", &buf)
 	if err != nil {
 		log.Printf("error: %s", err)
 		die()
 	}
 
-	_, err = io.Copy(out, resp.Body)
-	if err != nil {
-		log.Printf("error writing profile: %s", err)
-		die()
-	}
-
 	log.Printf("finished")
-
 }
 
 func doControLVersionCmd(cmd *cobra.Command, args []string) {

cmd/control_pprof.go

@@ -0,0 +1,80 @@
+package cmd
+
+import (
+	"net/http"
+	// FIXME: importing this package has the side-effect of poisoning the http.DefaultServeMux
+	// FIXME: with the /debug/pprof endpoints
+	"context"
+	"net"
+	"net/http/pprof"
+)
+
+type PProfServer struct {
+	cc       chan PprofServerControlMsg
+	state    PprofServerControlMsg
+	listener net.Listener
+}
+
+type PprofServerControlMsg struct {
+	// Whether the server should listen for requests on the given address
+	Run bool
+	// Must be set if Run is true, undefined otherwise
+	HttpListenAddress string
+}
+
+func NewPProfServer(ctx context.Context) *PProfServer {
+
+	s := &PProfServer{
+		cc: make(chan PprofServerControlMsg),
+	}
+
+	go s.controlLoop(ctx)
+	return s
+}
+
+func (s *PProfServer) controlLoop(ctx context.Context) {
+outer:
+	for {
+
+		var msg PprofServerControlMsg
+		select {
+		case <-ctx.Done():
+			if s.listener != nil {
+				s.listener.Close()
+			}
+			break outer
+		case msg = <-s.cc:
+			// proceed
+		}
+
+		var err error
+		if msg.Run && s.listener == nil {
+
+			s.listener, err = net.Listen("tcp", msg.HttpListenAddress)
+			if err != nil {
+				s.listener = nil
+				continue
+			}
+
+			// FIXME: because net/http/pprof does not provide a mux,
+			mux := http.NewServeMux()
+			mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
+			mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
+			mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
+			mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
+			mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
+			go http.Serve(s.listener, mux)
+			continue
+		}
+
+		if !msg.Run && s.listener != nil {
+			s.listener.Close()
+			s.listener = nil
+			continue
+		}
+	}
+}
+
+func (s *PProfServer) Control(msg PprofServerControlMsg) {
+	s.cc <- msg
+}