feat: add cpu and memory profiling (#504)

Signed-off-by: Mikhail Swift <mikhail@testifysec.com>

cmd/root.go

@@ -17,6 +17,8 @@ package cmd
 import (
 	"fmt"
 	"os"
+	"runtime"
+	"runtime/pprof"
 
 	"github.com/in-toto/go-witness/log"
 	_ "github.com/in-toto/go-witness/signer/kms/aws"
@@ -25,7 +27,10 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var ro = &options.RootOptions{}
+var (
+	ro             = &options.RootOptions{}
+	cpuProfileFile *os.File
+)
 
 func New() *cobra.Command {
 	cmd := &cobra.Command{
@@ -46,6 +51,7 @@ func New() *cobra.Command {
 	cmd.AddCommand(versionCmd())
 	cmd.AddCommand(AttestorsCmd())
 	cobra.OnInitialize(func() { preRoot(cmd, ro, logger) })
+	cobra.OnFinalize((func() { postRoot(ro, logger) }))
 	return cmd
 }
 
@@ -64,6 +70,40 @@ func preRoot(cmd *cobra.Command, ro *options.RootOptions, logger *logrusLogger)
 	if err := initConfig(cmd, ro); err != nil {
 		logger.l.Fatal(err)
 	}
+
+	var err error
+	if len(ro.CpuProfileFile) > 0 {
+		cpuProfileFile, err = os.Create(ro.CpuProfileFile)
+		if err != nil {
+			logger.l.Fatalf("could not create CPU profile: %v", err)
+		}
+
+		if err = pprof.StartCPUProfile(cpuProfileFile); err != nil {
+			logger.l.Fatalf("could not start CPU profile: %v", err)
+		}
+	}
+}
+
+func postRoot(ro *options.RootOptions, logger *logrusLogger) {
+	if cpuProfileFile != nil {
+		pprof.StopCPUProfile()
+		if err := cpuProfileFile.Close(); err != nil {
+			logger.l.Fatalf("could not close cpu profile file: %v", err)
+		}
+	}
+
+	if len(ro.MemProfileFile) > 0 {
+		memProfileFile, err := os.Create(ro.MemProfileFile)
+		if err != nil {
+			logger.l.Fatalf("could not create memory profile file: %v", err)
+		}
+
+		defer memProfileFile.Close()
+		runtime.GC()
+		if err := pprof.WriteHeapProfile(memProfileFile); err != nil {
+			logger.l.Fatalf("could not write memory profile: %v", err)
+		}
+	}
 }
 
 func loadOutfile(outFilePath string) (*os.File, error) {

docs/commands.md

@@ -19,8 +19,10 @@ Get information about all the available attestors in Witness
 ### Options inherited from parent commands
 
 ```
-  -c, --config string      Path to the witness config file (default ".witness.yaml")
-  -l, --log-level string   Level of logging to output (debug, info, warn, error) (default "info")
+  -c, --config string                   Path to the witness config file (default ".witness.yaml")
+      --debug-cpu-profile-file string   Path to store the CPU profile. Profiling will be enabled if this is non-empty
+      --debug-mem-profile-file string   Path to store the Memory profile. Profiling will be enabled if this is non-empty
+  -l, --log-level string                Level of logging to output (debug, info, warn, error) (default "info")
 ```
 
 ### SEE ALSO
@@ -88,8 +90,10 @@ witness run [cmd] [flags]
 ### Options inherited from parent commands
 
 ```
-  -c, --config string      Path to the witness config file (default ".witness.yaml")
-  -l, --log-level string   Level of logging to output (debug, info, warn, error) (default "info")
+  -c, --config string                   Path to the witness config file (default ".witness.yaml")
+      --debug-cpu-profile-file string   Path to store the CPU profile. Profiling will be enabled if this is non-empty
+      --debug-mem-profile-file string   Path to store the Memory profile. Profiling will be enabled if this is non-empty
+  -l, --log-level string                Level of logging to output (debug, info, warn, error) (default "info")
 ```
 
 ### SEE ALSO
@@ -148,8 +152,10 @@ witness sign [file] [flags]
 ### Options inherited from parent commands
 
 ```
-  -c, --config string      Path to the witness config file (default ".witness.yaml")
-  -l, --log-level string   Level of logging to output (debug, info, warn, error) (default "info")
+  -c, --config string                   Path to the witness config file (default ".witness.yaml")
+      --debug-cpu-profile-file string   Path to store the CPU profile. Profiling will be enabled if this is non-empty
+      --debug-mem-profile-file string   Path to store the Memory profile. Profiling will be enabled if this is non-empty
+  -l, --log-level string                Level of logging to output (debug, info, warn, error) (default "info")
 ```
 
 ### SEE ALSO
@@ -208,8 +214,10 @@ witness verify [flags]
 ### Options inherited from parent commands
 
 ```
-  -c, --config string      Path to the witness config file (default ".witness.yaml")
-  -l, --log-level string   Level of logging to output (debug, info, warn, error) (default "info")
+  -c, --config string                   Path to the witness config file (default ".witness.yaml")
+      --debug-cpu-profile-file string   Path to store the CPU profile. Profiling will be enabled if this is non-empty
+      --debug-mem-profile-file string   Path to store the Memory profile. Profiling will be enabled if this is non-empty
+  -l, --log-level string                Level of logging to output (debug, info, warn, error) (default "info")
 ```
 
 ### SEE ALSO
@@ -237,8 +245,10 @@ witness version [flags]
 ### Options inherited from parent commands
 
 ```
-  -c, --config string      Path to the witness config file (default ".witness.yaml")
-  -l, --log-level string   Level of logging to output (debug, info, warn, error) (default "info")
+  -c, --config string                   Path to the witness config file (default ".witness.yaml")
+      --debug-cpu-profile-file string   Path to store the CPU profile. Profiling will be enabled if this is non-empty
+      --debug-mem-profile-file string   Path to store the Memory profile. Profiling will be enabled if this is non-empty
+  -l, --log-level string                Level of logging to output (debug, info, warn, error) (default "info")
 ```
 
 ### SEE ALSO

options/root.go

@@ -14,14 +14,20 @@
 
 package options
 
-import "github.com/spf13/cobra"
+import (
+	"github.com/spf13/cobra"
+)
 
 type RootOptions struct {
-	Config   string
-	LogLevel string
+	Config         string
+	LogLevel       string
+	CpuProfileFile string
+	MemProfileFile string
 }
 
 func (ro *RootOptions) AddFlags(cmd *cobra.Command) {
 	cmd.PersistentFlags().StringVarP(&ro.Config, "config", "c", ".witness.yaml", "Path to the witness config file")
 	cmd.PersistentFlags().StringVarP(&ro.LogLevel, "log-level", "l", "info", "Level of logging to output (debug, info, warn, error)")
+	cmd.PersistentFlags().StringVar(&ro.CpuProfileFile, "debug-cpu-profile-file", "", "Path to store the CPU profile. Profiling will be enabled if this is non-empty")
+	cmd.PersistentFlags().StringVar(&ro.MemProfileFile, "debug-mem-profile-file", "", "Path to store the Memory profile. Profiling will be enabled if this is non-empty")
 }