cmd: introduce control socket & subcommand

Move pprof debugging there.

cmd/config.go

@@ -27,6 +27,9 @@ type Global struct {
 			SockDir string
 		}
 	}
+	Control struct {
+		Sockpath string
+	}
 }
 
 type JobDebugSettings struct {

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

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 {

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

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
-
 }