From f992fed968d3c691c5ceb3a39dc4876892305b2e Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Sat, 17 Feb 2018 16:20:10 +0100 Subject: [PATCH] control pprof rewrite: expose pprof metrics via HTTP server controlled from CLI --- cmd/config_job_control.go | 16 ++++++-- cmd/control.go | 75 ++++++++++++++++-------------------- cmd/control_pprof.go | 80 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 126 insertions(+), 45 deletions(-) create mode 100644 cmd/control_pprof.go diff --git a/cmd/config_job_control.go b/cmd/config_job_control.go index 1e3c675..dc8a1e4 100644 --- a/cmd/config_job_control.go +++ b/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 diff --git a/cmd/control.go b/cmd/control.go index 7c57dd7..3c9e7df 100644 --- a/cmd/control.go +++ b/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) { diff --git a/cmd/control_pprof.go b/cmd/control_pprof.go new file mode 100644 index 0000000..a6b6939 --- /dev/null +++ b/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 +}