From 12b9c7a346d74e723a08ef37c9b9e7e36c175da2 Mon Sep 17 00:00:00 2001
From: Richard Case <richard.case@outlook.com>
Date: Thu, 18 Aug 2022 14:47:31 +0100
Subject: [PATCH] feat: add pprof to flintlock

This change adds the ability to make pprof available via http. By
default it won't be available. If you specify a debug endpoint via
`--debug-endpoint` then it will start a webserver listening at that
endpoint and pprof is available from there.

We decided to use an endpoint as that will allow users to only expose
the debug endpoint locally or via the network.

Signed-off-by: Richard Case <richard.case@outlook.com>
---
 internal/command/flags/flags.go |  8 ++++++
 internal/command/run/run.go     | 46 +++++++++++++++++++++++++++++++++
 internal/config/config.go       |  4 ++-
 3 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/internal/command/flags/flags.go b/internal/command/flags/flags.go
index 5deaf597..860b6648 100644
--- a/internal/command/flags/flags.go
+++ b/internal/command/flags/flags.go
@@ -27,6 +27,7 @@ const (
 	tlsKeyFlag            = "tls-key"
 	tlsClientValidateFlag = "tls-client-validate"
 	tlsClientCAFlag       = "tls-client-ca"
+	debugEndpointFlag     = "debug-endpoint"
 )
 
 // AddGRPCServerFlagsToCommand will add gRPC server flags to the supplied command.
@@ -178,3 +179,10 @@ func AddContainerDFlagsToCommand(cmd *cobra.Command, cfg *config.Config) {
 		defaults.ContainerdNamespace,
 		"The name of the containerd namespace to use.")
 }
+
+func AddDebugFlagsToCommand(cmd *cobra.Command, cfg *config.Config) {
+	cmd.Flags().StringVar(&cfg.DebugEndpoint,
+		debugEndpointFlag,
+		"",
+		"The endpoint for the debug web server to listen on. It must include a port (e.g. localhost:10500).  An empty string means disable the debug endpoint.")
+}
diff --git a/internal/command/run/run.go b/internal/command/run/run.go
index 2c0e7a0a..167ead9e 100644
--- a/internal/command/run/run.go
+++ b/internal/command/run/run.go
@@ -9,6 +9,9 @@ import (
 	"os"
 	"os/signal"
 	"sync"
+	"time"
+
+	_ "net/http/pprof"
 
 	grpc_mw "github.com/grpc-ecosystem/go-grpc-middleware"
 	grpc_auth "github.com/grpc-ecosystem/go-grpc-middleware/auth"
@@ -59,6 +62,7 @@ func NewCommand(cfg *config.Config) (*cobra.Command, error) {
 	cmdflags.AddTLSFlagsToCommand(cmd, cfg)
 	cmdflags.AddContainerDFlagsToCommand(cmd, cfg)
 	cmdflags.AddFirecrackerFlagsToCommand(cmd, cfg)
+	cmdflags.AddDebugFlagsToCommand(cmd, cfg)
 
 	if err := cmdflags.AddNetworkFlagsToCommand(cmd, cfg); err != nil {
 		return nil, fmt.Errorf("adding network flags to run command: %w", err)
@@ -81,6 +85,18 @@ func runServer(ctx context.Context, cfg *config.Config) error {
 	wg := &sync.WaitGroup{}
 	ctx, cancel := context.WithCancel(log.WithLogger(ctx, logger))
 
+	if cfg.DebugEndpoint != "" {
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+
+			if err := runPProf(ctx, cfg); err != nil {
+				logger.Errorf("failed serving api: %v", err)
+			}
+		}()
+	}
+
 	if !cfg.DisableAPI {
 		wg.Add(1)
 
@@ -225,3 +241,33 @@ func generateOpts(ctx context.Context, cfg *config.Config) ([]grpc.ServerOption,
 
 	return opts, nil
 }
+
+func runPProf(ctx context.Context, cfg *config.Config) error {
+	logger := log.GetLogger(ctx)
+	logger.Warnf("Debug endpoint is ENABLED at %s", cfg.DebugEndpoint)
+
+	srv := &http.Server{
+		Addr:    cfg.DebugEndpoint,
+		Handler: http.DefaultServeMux,
+	}
+
+	go func() {
+		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			logger.Fatalf("starting debug endpoint: %v", err)
+		}
+	}()
+
+	<-ctx.Done()
+	logger.Debug("Exiting")
+
+	shutDownCtx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer func() {
+		cancel()
+	}()
+
+	if err := srv.Shutdown(shutDownCtx); err != nil {
+		logger.Warnf("Debug server shutdown failed:%+v", err)
+	}
+
+	return nil
+}
diff --git a/internal/config/config.go b/internal/config/config.go
index a2c91f30..021fa07d 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -14,7 +14,7 @@ type Config struct {
 	Logging log.Config
 	// GRPCEndpoint is the endpoint for the gRPC server.
 	GRPCAPIEndpoint string
-	// HTTPAPIEndpoint is the endpoint for the HTTP proxy for the gRPC service..
+	// HTTPAPIEndpoint is the endpoint for the HTTP proxy for the gRPC service
 	HTTPAPIEndpoint string
 	// FirecrackerBin is the firecracker binary to use.
 	FirecrackerBin string
@@ -46,6 +46,8 @@ type Config struct {
 	BasicAuthToken string
 	// TLS holds the TLS related configuration.
 	TLS TLSConfig
+	// DebugEndpoint is the endpoint for the debug web server. An empty string means disable the debug endpoint.
+	DebugEndpoint string
 }
 
 // TLSConfig holds the configuration for TLS.