kubectl: Refactor profiling helpers

Author: tchap

staging/src/k8s.io/kubectl/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()

staging/src/k8s.io/kubectl/pkg/cmd/profiling.go

@@ -18,6 +18,7 @@ package cmd
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"os/signal"
 	"runtime"
@@ -36,22 +37,25 @@ func addProfilingFlags(flags *pflag.FlagSet) {
 	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.
@@ -62,25 +66,29 @@ func initProfiling() error {
 	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
@@ -94,10 +102,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)
 	}