diff --git a/cmd/config.go b/cmd/config.go index a9fb160..09ad7cf 100644 --- a/cmd/config.go +++ b/cmd/config.go @@ -27,6 +27,9 @@ type Global struct { SockDir string } } + Control struct { + Sockpath string + } } type JobDebugSettings struct { diff --git a/cmd/config_job_control.go b/cmd/config_job_control.go new file mode 100644 index 0000000..0c7087b --- /dev/null +++ b/cmd/config_job_control.go @@ -0,0 +1,88 @@ +package cmd + +import ( + "context" + "fmt" + "github.com/pkg/errors" + "github.com/zrepl/zrepl/util" + "net" + "net/http" + "net/http/pprof" +) + +type ControlJob struct { + Name string + sockaddr *net.UnixAddr +} + +func NewControlJob(name, sockpath string) (j *ControlJob, err error) { + j = &ControlJob{Name: name} + + j.sockaddr, err = net.ResolveUnixAddr("unix", sockpath) + if err != nil { + err = errors.Wrap(err, "cannot resolve unix address") + return + } + + return +} + +func (j *ControlJob) JobName() string { + return j.Name +} + +const ( + ControlJobEndpointProfile string = "/debug/pprof/profile" +) + +func (j *ControlJob) JobStart(ctx context.Context) { + + log := ctx.Value(contextKeyLog).(Logger) + defer log.Printf("control job finished") + + l, err := ListenUnixPrivate(j.sockaddr) + if err != nil { + log.Printf("error listening: %s", err) + return + } + + mux := http.NewServeMux() + mux.Handle(ControlJobEndpointProfile, requestLogger{log, pprof.Profile}) + server := http.Server{Handler: mux} + +outer: + for { + + served := make(chan error) + go func() { + served <- server.Serve(l) + close(served) + }() + + select { + case <-ctx.Done(): + log.Printf("context: %s", ctx.Err()) + server.Shutdown(context.Background()) + break outer + case err = <-served: + if err != nil { + log.Printf("error serving: %s", err) + break outer + } + } + + } + +} + +type requestLogger struct { + log Logger + handlerFunc func(w http.ResponseWriter, r *http.Request) +} + +func (l requestLogger) ServeHTTP(w http.ResponseWriter, r *http.Request) { + log := util.NewPrefixLogger(l.log, fmt.Sprintf("%s %s", r.Method, r.URL)) + log.Printf("start") + l.handlerFunc(w, r) + log.Printf("finish") +} diff --git a/cmd/config_parse.go b/cmd/config_parse.go index 6864105..65c754b 100644 --- a/cmd/config_parse.go +++ b/cmd/config_parse.go @@ -16,6 +16,14 @@ var ConfigFileDefaultLocations []string = []string{ "/usr/local/etc/zrepl/zrepl.yml", } +const ( + JobNameControl string = "control" +) + +var ReservedJobNames []string = []string{ + JobNameControl, +} + type ConfigParsingContext struct { Global *Global } @@ -73,6 +81,8 @@ func parseConfig(i interface{}) (c *Config, err error) { // Parse global with defaults c.Global.Serve.Stdinserver.SockDir = "/var/run/zrepl/stdinserver" + c.Global.Control.Sockpath = "/var/run/zrepl/control" + err = mapstructure.Decode(asMap.Global, &c.Global) if err != nil { err = errors.Wrap(err, "cannot parse global section: %s") @@ -82,7 +92,7 @@ func parseConfig(i interface{}) (c *Config, err error) { cpc := ConfigParsingContext{&c.Global} jpc := JobParsingContext{cpc} - // Parse Jobs + // Jobs c.Jobs = make(map[string]Job, len(asMap.Jobs)) for i := range asMap.Jobs { job, err := parseJob(jpc, asMap.Jobs[i]) @@ -95,9 +105,21 @@ func parseConfig(i interface{}) (c *Config, err error) { err = errors.Wrapf(err, "cannot parse job '%v'", namei) return nil, err } + jn := job.JobName() + if _, ok := c.Jobs[jn]; ok { + err = errors.Errorf("duplicate job name: %s", jn) + return nil, err + } c.Jobs[job.JobName()] = job } + cj, err := NewControlJob(JobNameControl, jpc.Global.Control.Sockpath) + if err != nil { + err = errors.Wrap(err, "cannot create control job") + return + } + c.Jobs[JobNameControl] = cj + return c, nil } @@ -131,10 +153,16 @@ func parseJob(c JobParsingContext, i map[string]interface{}) (j Job, err error) return } + for _, r := range ReservedJobNames { + if name == r { + err = errors.Errorf("job name '%s' is reserved", name) + return + } + } + jobtype, err := extractStringField(i, "type", true) if err != nil { return - } switch jobtype { diff --git a/cmd/control.go b/cmd/control.go new file mode 100644 index 0000000..dfdd817 --- /dev/null +++ b/cmd/control.go @@ -0,0 +1,101 @@ +package cmd + +import ( + "context" + "fmt" + "github.com/spf13/cobra" + "io" + golog "log" + "net" + "net/http" + "net/url" + "os" +) + +var controlCmd = &cobra.Command{ + Use: "control", + Short: "control zrepl daemon", +} + +var pprofCmd = &cobra.Command{ + Use: "pprof cpu OUTFILE", + Short: "pprof CPU of daemon to OUTFILE (- for stdout)", + Run: doControlPProf, +} +var pprofCmdArgs struct { + seconds int64 +} + +func init() { + RootCmd.AddCommand(controlCmd) + controlCmd.AddCommand(pprofCmd) + pprofCmd.Flags().Int64Var(&pprofCmdArgs.seconds, "seconds", 30, "seconds to profile") +} + +func doControlPProf(cmd *cobra.Command, args []string) { + + log := golog.New(os.Stderr, "", 0) + + die := func() { + log.Printf("exiting after error") + os.Exit(1) + } + + ctx := context.WithValue(context.Background(), contextKeyLog, log) + conf, err := ParseConfig(ctx, rootArgs.configFile) + if err != nil { + log.Printf("error parsing config: %s", err) + die() + } + + 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 + 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") + httpc := http.Client{ + Transport: &http.Transport{ + DialContext: func(_ context.Context, _, _ string) (net.Conn, error) { + return net.Dial("unix", conf.Global.Control.Sockpath) + }, + }, + } + + 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()) + 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") + +} diff --git a/cmd/main.go b/cmd/main.go index 1dde372..5c3e331 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -12,8 +12,6 @@ package cmd import ( "github.com/spf13/cobra" - "net/http" - _ "net/http/pprof" ) type Logger interface { @@ -33,24 +31,9 @@ var RootCmd = &cobra.Command{ var rootArgs struct { configFile string - httpPprof string } func init() { //cobra.OnInitialize(initConfig) RootCmd.PersistentFlags().StringVar(&rootArgs.configFile, "config", "", "config file path") - RootCmd.PersistentFlags().StringVar(&rootArgs.httpPprof, "debug.pprof.http", "", "run pprof http server on given port") -} - -func initConfig() { - - // CPU profiling - if rootArgs.httpPprof != "" { - go func() { - http.ListenAndServe(rootArgs.httpPprof, nil) - }() - } - - return - }