cmd/cue: add global cpuprofile and memprofile flags

Mirroring those available via `go help testflag`.
These are helpful to get an idea of why a cmd/cue command is slow,
for example `cue fmt --check ./...` in the CUE repo taking over 10s.

The flags are global to be available for all commands, but hidden,
as they should only be necessary for debugging performance issues.
We can choose to un-hide them at a later point.

Add a small regression test as well, with some basic checks
that the pprof files get written and are minimally valid.

Signed-off-by: Daniel Martí <mvdan@mvdan.cc>
Change-Id: I9491fe67c6f3f8657cfa586a76fcd7bb46b3d429
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/1192973
TryBot-Result: CUEcueckoo <cueckoo@cuelang.org>
Reviewed-by: Noam Dolovich <noam.tzvi.dolovich@gmail.com>
Reviewed-by: Roger Peppe <rogpeppe@gmail.com>
Unity-Result: CUE porcuepine <cue.porcuepine@gmail.com>

cmd/cue/cmd/flags.go

@@ -50,6 +50,10 @@ const (
 	flagWithContext flagName = "with-context"
 	flagOut         flagName = "out"
 	flagOutFile     flagName = "outfile"
+
+	// Hidden flags.
+	flagCpuProfile flagName = "cpuprofile"
+	flagMemProfile flagName = "memprofile"
 )
 
 func addOutFlags(f *pflag.FlagSet, allowNonCUE bool) {
@@ -74,6 +78,11 @@ func addGlobalFlags(f *pflag.FlagSet) {
 	f.BoolP(string(flagVerbose), "v", false,
 		"print information about progress")
 	f.BoolP(string(flagAllErrors), "E", false, "print all available errors")
+
+	f.String(string(flagCpuProfile), "", "write a CPU profile to the specified file before exiting")
+	f.MarkHidden(string(flagCpuProfile))
+	f.String(string(flagMemProfile), "", "write an allocation profile to the specified file before exiting")
+	f.MarkHidden(string(flagMemProfile))
 }
 
 func addOrphanFlags(f *pflag.FlagSet) {

cmd/cue/cmd/root.go

@@ -17,9 +17,11 @@ package cmd
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
 	"os"
 	"runtime"
+	"runtime/pprof"
 	"testing"
 
 	"github.com/spf13/cobra"
@@ -96,8 +98,33 @@ func mkRunE(c *Command, f runFunction) func(*cobra.Command, []string) error {
 		// We don't want that work to count towards $CUE_STATS.
 		adt.ResetStats()
 
+		if cpuprofile := flagCpuProfile.String(c); cpuprofile != "" {
+			f, err := os.Create(cpuprofile)
+			if err != nil {
+				return fmt.Errorf("could not create CPU profile: %v", err)
+			}
+			defer f.Close()
+			if err := pprof.StartCPUProfile(f); err != nil {
+				return fmt.Errorf("could not start CPU profile: %v", err)
+			}
+			defer pprof.StopCPUProfile()
+		}
+
 		err := f(c, args)
 
+		// TODO(mvdan): support -memprofilerate like `go help testflag`.
+		if memprofile := flagMemProfile.String(c); memprofile != "" {
+			f, err := os.Create(memprofile)
+			if err != nil {
+				return fmt.Errorf("could not create memory profile: %v", err)
+			}
+			defer f.Close()
+			runtime.GC() // get up-to-date statistics
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				return fmt.Errorf("could not write memory profile: %v", err)
+			}
+		}
+
 		if statsEnc != nil {
 			var stats Stats
 			stats.CUE = adt.TotalStats()

cmd/cue/cmd/testdata/script/pprof.txtar

@@ -0,0 +1,20 @@
+# Grab a CPU profile and check that it's valid.
+exec cue eval --cpuprofile cpu.out x.cue
+go tool pprof -top cpu.out
+stdout 'Type: cpu'
+
+# Grab an allocation profile and check that it's valid.
+exec cue eval --memprofile mem.out x.cue
+go tool pprof -top mem.out
+stdout 'Type: inuse_space'
+
+# Both flags can be used at the same time.
+rm cpu.out mem.out
+exec cue eval --cpuprofile=cpu.out --memprofile=mem.out x.cue
+exists cpu.out
+exists mem.out
+
+-- x.cue --
+a: 1
+b: 2
+c: a | b