cmd/uplink: add options to save pprof/trace information

Change-Id: I5bcc602366de4ebd9b761e641a3806ddaeb9ecba

cmd/uplink/external.go

@@ -12,6 +12,8 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"runtime/pprof"
+	"runtime/trace"
 	"strconv"
 	"strings"
 	"sync"
@@ -63,6 +65,11 @@ type external struct {
 		verbose      bool    // flag to print out tracing information (like the used trace id)
 	}
 
+	debug struct {
+		pprofFile string
+		traceFile string
+	}
+
 	events struct {
 		address string // if non-zero, events are sent to this address.
 	}
@@ -125,6 +132,16 @@ func (ex *external) Setup(f clingy.Flags) {
 		clingy.Advanced,
 	).(string)
 
+	ex.debug.pprofFile = f.Flag(
+		"debug-pprof", "File to collect Golang pprof profiling data", "",
+		clingy.Advanced,
+	).(string)
+
+	ex.debug.traceFile = f.Flag(
+		"debug-trace", "File to collect Golang trace data", "",
+		clingy.Advanced,
+	).(string)
+
 	ex.dirs.loaded = true
 }
 
@@ -205,14 +222,14 @@ func (ex *external) analyticsEnabled() bool {
 
 // Wrap is called by clingy with the command to be executed.
 func (ex *external) Wrap(ctx context.Context, cmd clingy.Command) (err error) {
-	if err := ex.migrate(); err != nil {
+	if err = ex.migrate(); err != nil {
 		return err
 	}
-	if err := ex.loadConfig(); err != nil {
+	if err = ex.loadConfig(); err != nil {
 		return err
 	}
 	if !ex.config.loaded {
-		if err := saveInitialConfig(ctx, ex); err != nil {
+		if err = saveInitialConfig(ctx, ex); err != nil {
 			return err
 		}
 	}
@@ -222,6 +239,40 @@ func (ex *external) Wrap(ctx context.Context, cmd clingy.Command) (err error) {
 		ctx = experiment.With(ctx, exp)
 	}
 
+	if ex.debug.pprofFile != "" {
+		var output *os.File
+		output, err = os.Create(ex.debug.pprofFile)
+		if err != nil {
+			return errs.Wrap(err)
+		}
+		defer func() {
+			err = errs.Combine(err, output.Close())
+		}()
+
+		err = pprof.StartCPUProfile(output)
+		if err != nil {
+			return errs.Wrap(err)
+		}
+		defer pprof.StopCPUProfile()
+	}
+
+	if ex.debug.traceFile != "" {
+		var output *os.File
+		output, err = os.Create(ex.debug.traceFile)
+		if err != nil {
+			return errs.Wrap(err)
+		}
+		defer func() {
+			err = errs.Combine(err, output.Close())
+		}()
+
+		err = trace.Start(output)
+		if err != nil {
+			return errs.Wrap(err)
+		}
+		defer trace.Stop()
+	}
+
 	// N.B.: Tracing is currently disabled by default (sample == 0, traceID == 0) and is
 	// something a user can only opt into. as a result, we don't check ex.analyticsEnabled()
 	// in this if statement. If we do ever start turning on trace samples by default, we