add profiling to backup/restore

cmd/backup.go

@@ -29,6 +29,7 @@ Backs up a FeatureBase server to a local, tar-formatted snapshot file.
 	flags.StringVar(&cmd.Host, "host", "localhost:10101", "The address (host:port) of FeatureBase (HTTP).")
 	flags.StringVar(&cmd.Index, "index", "", "Index to backup, default backs up all indexes. ")
 	flags.DurationVar(&cmd.RetryPeriod, "retry-period", cmd.RetryPeriod, "Length of time after HTTP request failure to continue retrying request.")
+	flags.StringVar(&cmd.Pprof, "pprof", cmd.Pprof, "host:port to listen for profiling requests at /debug/pprof and /debug/fgprof.")
 	ctl.SetTLSConfig(flags, "", &cmd.TLS.CertificatePath, &cmd.TLS.CertificateKeyPath, &cmd.TLS.CACertPath, &cmd.TLS.SkipVerify, &cmd.TLS.EnableClientVerification)
 	return ccmd
 }

cmd/restore.go

@@ -26,6 +26,7 @@ The Restore command will take a backup archive and restore it to a new, clean cl
 	flags.StringVar(&cmd.Host, "host", "localhost:10101", "host:port of FeatureBase.")
 	flags.IntVar(&cmd.Concurrency, "concurrency", 1, "number of concurrent uploads")
 	flags.DurationVar(&cmd.RetryPeriod, "retry-period", cmd.RetryPeriod, "Length of time after HTTP request failure to continue retrying request.")
+	flags.StringVar(&cmd.Pprof, "pprof", cmd.Pprof, "host:port to listen for profiling requests at /debug/pprof and /debug/fgprof.")
 	ctl.SetTLSConfig(
 		flags, "",
 		&cmd.TLS.CertificatePath,

ctl/backup.go

@@ -13,9 +13,10 @@ import (
 	"time"
 
 	pilosa "github.com/molecula/featurebase/v2"
-	"github.com/molecula/featurebase/v2/http"
+	fb_http "github.com/molecula/featurebase/v2/http"
 	"github.com/molecula/featurebase/v2/server"
 	"github.com/molecula/featurebase/v2/topology"
+	"github.com/pkg/errors"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -41,6 +42,9 @@ type BackupCommand struct { // nolint: maligned
 	// Amount of time after first failed request to continue retrying.
 	RetryPeriod time.Duration `json:"retry-period"`
 
+	// Host:port on which to listen for pprof.
+	Pprof string `json:"pprof"`
+
 	// Reusable client.
 	client pilosa.InternalClient
 
@@ -56,11 +60,19 @@ func NewBackupCommand(stdin io.Reader, stdout, stderr io.Writer) *BackupCommand
 		CmdIO:       pilosa.NewCmdIO(stdin, stdout, stderr),
 		Concurrency: 1,
 		RetryPeriod: time.Minute,
+		Pprof:       "localhost:43809",
 	}
 }
 
 // Run executes the main program execution.
 func (cmd *BackupCommand) Run(ctx context.Context) (err error) {
+	logger := cmd.Logger()
+	close, err := startProfilingServer(cmd.Pprof, logger)
+	if err != nil {
+		return errors.Wrap(err, "starting profiling server")
+	}
+	defer close()
+
 	// Validate arguments.
 	if cmd.OutputDir == "" {
 		return fmt.Errorf("-o flag required")
@@ -75,7 +87,7 @@ func (cmd *BackupCommand) Run(ctx context.Context) (err error) {
 	}
 
 	// Create a client to the server.
-	client, err := commandClient(cmd, http.WithClientRetryPeriod(cmd.RetryPeriod))
+	client, err := commandClient(cmd, fb_http.WithClientRetryPeriod(cmd.RetryPeriod))
 	if err != nil {
 		return fmt.Errorf("creating client: %w", err)
 	}
@@ -267,7 +279,7 @@ func (cmd *BackupCommand) backupShardNode(ctx context.Context, indexName string,
 	logger := cmd.Logger()
 	logger.Printf("backing up shard: index=%q id=%d", indexName, shard)
 
-	client := http.NewInternalClientFromURI(&node.URI, http.GetHTTPClient(cmd.tlsConfig), http.WithClientRetryPeriod(cmd.RetryPeriod))
+	client := fb_http.NewInternalClientFromURI(&node.URI, fb_http.GetHTTPClient(cmd.tlsConfig), fb_http.WithClientRetryPeriod(cmd.RetryPeriod))
 	rc, err := client.ShardReader(ctx, indexName, shard)
 	if err != nil {
 		return fmt.Errorf("fetching shard reader: %w", err)

ctl/restore.go

@@ -5,7 +5,6 @@ import (
 	"context"
 	"crypto/tls"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -21,6 +20,7 @@ import (
 	fb_http "github.com/molecula/featurebase/v2/http"
 	"github.com/molecula/featurebase/v2/server"
 	"github.com/molecula/featurebase/v2/topology"
+	"github.com/pkg/errors"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -37,6 +37,9 @@ type RestoreCommand struct {
 	// Amount of time after first failed request to continue retrying.
 	RetryPeriod time.Duration `json:"retry-period"`
 
+	// Host:port on which to listen for pprof.
+	Pprof string `json:"pprof"`
+
 	// Reusable client.
 	client pilosa.InternalClient
 
@@ -51,12 +54,18 @@ func NewRestoreCommand(stdin io.Reader, stdout, stderr io.Writer) *RestoreComman
 		CmdIO:       pilosa.NewCmdIO(stdin, stdout, stderr),
 		RetryPeriod: time.Second * 30,
 		Concurrency: 1,
+		Pprof:       "localhost:43809",
 	}
 }
 
 // Run executes the restore.
 func (cmd *RestoreCommand) Run(ctx context.Context) (err error) {
 	logger := cmd.Logger()
+	close, err := startProfilingServer(cmd.Pprof, logger)
+	if err != nil {
+		return errors.Wrap(err, "starting profiling server")
+	}
+	defer close()
 
 	// Validate arguments.
 	if cmd.Path == "" {

ctl/util.go

@@ -0,0 +1,56 @@
+package ctl
+
+import (
+	"context"
+	"net"
+	"net/http"
+	"net/http/pprof"
+	"runtime"
+	"time"
+
+	"github.com/felixge/fgprof"
+	"github.com/molecula/featurebase/v2/logger"
+	"github.com/pkg/errors"
+)
+
+// startProfilingServer starts a server which handles /debug/pprof and
+// /debug/fgprof for use in utilities we might want to profile but
+// wouldn't otherwise be running an http server. Caller should call
+// the returned close function before exiting to release resources.
+func startProfilingServer(addr string, logger logger.Logger) (close func() error, err error) {
+	if addr == "" {
+		return func() error { return nil }, nil
+	}
+
+	sm := http.NewServeMux()
+	sm.Handle("/debug/fgprof", fgprof.Handler())
+	sm.HandleFunc("/debug/pprof/", pprof.Index)
+	sm.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	sm.HandleFunc("/debug/pprof/profile", pprof.Profile)
+	sm.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	sm.HandleFunc("/debug/pprof/trace", pprof.Trace)
+	s := &http.Server{
+		Addr:    addr,
+		Handler: sm,
+	}
+	runtime.SetBlockProfileRate(10000000) // 1 sample per 10 ms
+	runtime.SetMutexProfileFraction(100)  // 1% sampling
+	ln, err := net.Listen("tcp", addr)
+	if err != nil {
+		return nil, err
+	}
+	go func() {
+		logger.Printf("Listening for /debug/pprof/ and /debug/fgprof on '%s'", addr)
+		logger.Printf("%v", s.Serve(ln))
+	}()
+
+	return func() error {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+		defer cancel()
+		err := s.Shutdown(ctx)
+		if err != nil {
+			return errors.Wrap(err, "shutting down profiling server")
+		}
+		return s.Close()
+	}, nil
+}

http/handler.go

@@ -3348,7 +3348,7 @@ func (h *Handler) handlePostRestore(w http.ResponseWriter, r *http.Request) {
 	//validate shard for this node
 	err = h.api.RestoreShard(ctx, indexName, shard, r.Body)
 	if err != nil {
-		http.Error(w, fmt.Sprintf("failed to restore shared %v %v err:%v", indexName, shard, err), http.StatusBadRequest)
+		http.Error(w, fmt.Sprintf("failed to restore shard %v %v err:%v", indexName, shard, err), http.StatusBadRequest)
 		return
 	}