Merge pull request #134709 from tchap/kubectl-tracing

kubectl: Add support for tracing

Kubernetes-commit: 8cd57a9e6f07d2446910d493781557b6cd4d6b99

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()

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