Add flag to enable pprof endpoints

This commit modifies the server, scraper & worker to accept a new `--debug` flag which
adds the pprof endpoints when set. Including pprof isn't really necessary for a produciton
deployment. It just adds convenience when running the project locally and we can just
scrape and process our own profiles.

Signed-off-by: David Bond <[email protected]>

cmd/scrape/scrape.go

@@ -36,6 +36,7 @@ func Command() *cobra.Command {
 		frequency  time.Duration
 		app        string
 		mode       string
+		debug      bool
 	)
 
 	cmd := &cobra.Command{
@@ -89,7 +90,8 @@ func Command() *cobra.Command {
 			})
 			group.Go(func() error {
 				return server.Run(ctx, server.Config{
-					Port: port,
+					Debug: debug,
+					Port:  port,
 					Middleware: []server.Middleware{
 						logger.Middleware(logger.FromContext(ctx)),
 					},
@@ -108,6 +110,7 @@ func Command() *cobra.Command {
 	flags.DurationVarP(&duration, "duration", "d", time.Second*30, "How long to profile targets for")
 	flags.DurationVarP(&frequency, "frequency", "f", time.Minute, "Interval between scraping targets")
 	flags.StringVarP(&mode, "mode", "m", modeFile, "Mode to use for obtaining targets (file, kube, nomad, consul)")
+	flags.BoolVar(&debug, "debug", false, "Enable debug endpoints")
 
 	cmd.MarkFlagRequired("app")
 	cmd.MarkFlagRequired("sample-size")

cmd/server/server.go

@@ -18,6 +18,7 @@ func Command() *cobra.Command {
 		port           int
 		eventWriterURL string
 		blobStoreURL   string
+		debug          bool
 	)
 
 	cmd := &cobra.Command{
@@ -44,7 +45,8 @@ func Command() *cobra.Command {
 			defer closers.Close(ctx, blobs)
 
 			return server.Run(ctx, server.Config{
-				Port: port,
+				Debug: debug,
+				Port:  port,
 				Controllers: []server.Controller{
 					profile.NewHTTPController(blobs, writer),
 				},
@@ -59,6 +61,7 @@ func Command() *cobra.Command {
 	flags.IntVarP(&port, "port", "p", 8080, "Port to use for HTTP traffic")
 	flags.StringVar(&eventWriterURL, "event-writer-url", "", "The URL to use for writing to the event bus")
 	flags.StringVar(&blobStoreURL, "blob-store-url", "", "The URL to use for connecting to blob storage")
+	flags.BoolVar(&debug, "debug", false, "Enable debug endpoints")
 
 	cmd.MarkPersistentFlagRequired("blob-store-url")
 	cmd.MarkPersistentFlagRequired("event-writer-url")

cmd/worker/worker.go

@@ -21,6 +21,7 @@ func Command() *cobra.Command {
 		blobStoreURL   string
 		prune          string
 		port           int
+		debug          bool
 	)
 
 	cmd := &cobra.Command{
@@ -79,7 +80,8 @@ func Command() *cobra.Command {
 			})
 			group.Go(func() error {
 				return server.Run(ctx, server.Config{
-					Port: port,
+					Debug: debug,
+					Port:  port,
 					Middleware: []server.Middleware{
 						logger.Middleware(logger.FromContext(ctx)),
 					},
@@ -96,6 +98,7 @@ func Command() *cobra.Command {
 	flags.StringVar(&blobStoreURL, "blob-store-url", "", "The URL to use for connecting to blob storage")
 	flags.IntVarP(&port, "port", "p", 8081, "Port to use for HTTP traffic")
 	flags.StringVar(&prune, "prune", "", "Location of the configuration file for profile pruning")
+	flags.BoolVar(&debug, "debug", false, "Enable debug endpoints")
 
 	cmd.MarkPersistentFlagRequired("blob-store-url")
 	cmd.MarkPersistentFlagRequired("event-reader-url")

internal/server/server.go

@@ -21,6 +21,8 @@ type (
 		Controllers []Controller
 		// Middleware functions to invoke prior to request handlers.
 		Middleware []Middleware
+		// Enables debug endpoints for pprof.
+		Debug bool
 	}
 
 	// The Controller interface describes types that register HTTP request handlers.
@@ -43,7 +45,9 @@ func Run(ctx context.Context, config Config) error {
 		controller.Register(mux)
 	}
 
-	registerDebug(mux)
+	if config.Debug {
+		registerDebug(mux)
+	}
 
 	server := &http.Server{
 		Handler: mux,