diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index 0c582ff33..58523a143 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -306,6 +306,7 @@ func HandlePluginCommand(pluginHandler PluginHandler, cmdArgs []string, minArgs func NewKubectlCommand(o KubectlOptions) *cobra.Command { warningHandler := rest.NewWarningWriter(o.IOStreams.ErrOut, rest.WarningWriterOptions{Deduplicate: true, Color: term.AllowsColorOutput(o.IOStreams.ErrOut)}) warningsAsErrors := false + var finishProfiling func() error // Parent command to which all subcommands are added. cmds := &cobra.Command{ Use: "kubectl", @@ -327,11 +328,15 @@ func NewKubectlCommand(o KubectlOptions) *cobra.Command { plugin.SetupPluginCompletion(cmd, args) } - return initProfiling() + var err error + finishProfiling, err = initProfiling() + return err }, PersistentPostRunE: func(*cobra.Command, []string) error { - if err := flushProfiling(); err != nil { - return err + if finishProfiling != nil { + if err := finishProfiling(); err != nil { + return err + } } if warningsAsErrors { count := warningHandler.WarningCount() diff --git a/pkg/cmd/profiling.go b/pkg/cmd/profiling.go index 758d7a116..c54d1038a 100644 --- a/pkg/cmd/profiling.go +++ b/pkg/cmd/profiling.go @@ -18,10 +18,12 @@ package cmd import ( "fmt" + "io" "os" "os/signal" "runtime" "runtime/pprof" + "runtime/trace" "github.com/spf13/pflag" ) @@ -32,26 +34,29 @@ var ( ) func addProfilingFlags(flags *pflag.FlagSet) { - flags.StringVar(&profileName, "profile", "none", "Name of profile to capture. One of (none|cpu|heap|goroutine|threadcreate|block|mutex)") + flags.StringVar(&profileName, "profile", "none", "Name of profile to capture. One of (none|cpu|heap|goroutine|threadcreate|block|mutex|trace)") flags.StringVar(&profileOutput, "profile-output", "profile.pprof", "Name of the file to write the profile to") } -func initProfiling() error { +// initProfiling inits profiling and returns a function to be called on exit to flush and close. +func initProfiling() (func() error, error) { var ( f *os.File err error ) switch profileName { case "none": - return nil + return nil, nil case "cpu": f, err = os.Create(profileOutput) if err != nil { - return err + return nil, err } + err = pprof.StartCPUProfile(f) if err != nil { - return err + f.Close() //nolint:errcheck + return nil, err } // Block and mutex profiles need a call to Set{Block,Mutex}ProfileRate to // output anything. We choose to sample all events. @@ -59,33 +64,53 @@ func initProfiling() error { runtime.SetBlockProfileRate(1) case "mutex": runtime.SetMutexProfileFraction(1) + case "trace": + f, err = os.Create(profileOutput) + if err != nil { + return nil, err + } + + // Enable the CPU profiler. Samples will be captured in the execution trace. + // This is the same rate value as used in pprof.StartCPUProfile. + runtime.SetCPUProfileRate(100) + if err := trace.Start(f); err != nil { + f.Close() //nolint:errcheck + return nil, err + } default: // Check the profile name is valid. if profile := pprof.Lookup(profileName); profile == nil { - return fmt.Errorf("unknown profile '%s'", profileName) + return nil, fmt.Errorf("unknown profile '%s'", profileName) } } - // If the command is interrupted before the end (ctrl-c), flush the - // profiling files + // If the command is interrupted before the end (ctrl-c), flush the profiling files c := make(chan os.Signal, 1) signal.Notify(c, os.Interrupt) go func() { <-c - f.Close() - flushProfiling() + flushProfiling(f) //nolint:errcheck os.Exit(0) }() - return nil + return func() error { + return flushProfiling(f) + }, nil } -func flushProfiling() error { +func flushProfiling(output io.Closer) error { + if output != nil { + defer output.Close() //nolint:errcheck + } + switch profileName { case "none": return nil case "cpu": pprof.StopCPUProfile() + case "trace": + trace.Stop() + runtime.SetCPUProfileRate(0) case "heap": runtime.GC() fallthrough @@ -94,10 +119,12 @@ func flushProfiling() error { if profile == nil { return nil } + f, err := os.Create(profileOutput) if err != nil { return err } + defer f.Close() profile.WriteTo(f, 0) }