feat: add pprof to flintlock

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

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

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.