Support CPU profiling

go/mlr.go

@@ -4,26 +4,59 @@ import (
 	"fmt"
 	"os"
 	"runtime"
+	"runtime/pprof"
 
 	"miller/cli"
 	"miller/stream"
 )
 
-// ----------------------------------------------------------------
-// xxx to do: stdout/stderr w/ ternary on exitrc
-func usage() {
-	// xxx temp grammar
-	fmt.Fprintf(os.Stderr, "Usage: %s [options] {ifmt} {mapper} {ofmt} {filenames ...}\n",
-		os.Args[0])
-	fmt.Fprintf(os.Stderr, "If no file name is given, or if filename is \"-\", stdin is used.\n")
-	// stub
-	os.Exit(1)
-}
-
 // ----------------------------------------------------------------
 func main() {
 	runtime.GOMAXPROCS(4) // Seems reasonable these days
 
+	//  - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+	// CPU profiling
+	//
+	// We do this here not in the command-line parser since
+	// pprof.StopCPUProfile() needs to be called at the very end of everything.
+	// Putting this pprof logic into a go func running in parallel with main,
+	// and properly stopping the profile only when main ends via chan-sync,
+	// results in a zero-length pprof file.
+	//
+	// Usage:
+	// * mlr --cpuprofile cpu.pprof put -f example.mlr then nothing ~/tmp/huge > /dev/null
+	// * go tool pprof mlr cpu.pprof
+	//   top10
+	// * go tool pprof --pdf mlr cpu.pprof > mlr-call-graph.pdf
+
+	if len(os.Args) >= 3 && os.Args[1] == "--cpuprofile" {
+		profFilename := os.Args[2]
+		handle, err := os.Create(profFilename)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, os.Args[0], ": ", "Could not start CPU profile: ", err)
+			return
+		}
+		defer handle.Close()
+
+		if err := pprof.StartCPUProfile(handle); err != nil {
+			fmt.Fprintln(os.Stderr, os.Args[0], ": ", "Could not start CPU profile: ", err)
+			return
+		}
+		defer pprof.StopCPUProfile()
+
+		fmt.Fprintf(os.Stderr, "CPU profile started.\n")
+		defer fmt.Fprintf(os.Stderr, "CPU profile finished.\n")
+		//	c := make(chan os.Signal)
+		//	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+		//	go func() {
+		//		<-c
+		//		pprof.StopCPUProfile()
+		//		os.Exit(0)
+		//	}()
+	}
+	//  - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+	// Start of Miller main per se
 	options, recordMappers, filenames, err := cli.ParseCommandLine(os.Args)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, os.Args[0], ": ", err)

go/src/miller/cli/mlrcli.go

@@ -63,13 +63,18 @@ func ParseCommandLine(args []string) (
 	for argi < argc /* variable increment: 1 or 2 depending on flag */ {
 		if args[argi][0] != '-' {
 			break // No more flag options to process
+		} else if args[argi] == "--cpuprofile" {
+			// Already handled in main(); ignore here.
+			checkArgCount(args, argi, argc, 1)
+			argi += 2
 		} else if handleTerminalUsage(args, argc, argi) {
 			os.Exit(0)
 		} else if handleReaderOptions(args, argc, &argi, &options.ReaderOptions) {
 			// handled
 		} else if handleWriterOptions(args, argc, &argi, &options.WriterOptions) {
 			// handled
-		} else if handleReaderWriterOptions(args, argc, &argi, &options.ReaderOptions, &options.WriterOptions) {
+		} else if handleReaderWriterOptions(args, argc, &argi,
+			&options.ReaderOptions, &options.WriterOptions) {
 			// handled
 		} else if handleMiscOptions(args, argc, &argi, &options) {
 			// handled
@@ -151,7 +156,13 @@ func ParseCommandLine(args []string) (
 // Returns a list of mappers, from the starting point in args given by *pargi.
 // Bumps *pargi to point to remaining post-mapper-setup args, i.e. filenames.
 
-func parseMappers(args []string, pargi *int, argc int, options *clitypes.TOptions) ([]mapping.IRecordMapper, error) {
+func parseMappers(
+	args []string,
+	pargi *int,
+	argc int,
+	options *clitypes.TOptions,
+) ([]mapping.IRecordMapper, error) {
+
 	mapperList := make([]mapping.IRecordMapper, 0)
 	argi := *pargi
 
@@ -2413,21 +2424,3 @@ func handleMiscOptions(
 //	MLR_INTERNAL_CODING_ERROR_UNLESS(lhmsll_has_key(pmap, key));
 //	return lhmsll_get(pmap, key);
 //}
-
-//	cpuprofile := flag.String("cpuprofile", "", "Write CPU profile to `file`")
-//// ----------------------------------------------------------------
-//func maybeProfile(cpuprofile *string) {
-//	// to do: move to method
-//	// go tool pprof mlr foo.prof
-//	//   top10
-//	if *cpuprofile != "" {
-//		f, err := os.Create(*cpuprofile)
-//		if err != nil {
-//			fmt.Fprintln(os.Stderr, os.Args[0], ": ", "Could not start CPU profile: ", err)
-//		}
-//		defer f.Close()
-//		if err := pprof.StartCPUProfile(f); err != nil {
-//			fmt.Fprintln(os.Stderr, os.Args[0], ": ", "Could not start CPU profile: ", err)
-//		}
-//		defer pprof.StopCPUProfile()
-//	}