diff --git a/cmd/root.go b/cmd/root.go index 81c15d51..ce86e013 100644 --- a/cmd/root.go +++ b/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) { diff --git a/docs/commands.md b/docs/commands.md index 58a07fa9..2de90877 100644 --- a/docs/commands.md +++ b/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 diff --git a/options/root.go b/options/root.go index 5a410845..8e84dd38 100644 --- a/options/root.go +++ b/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") }