pkg/cmdprof: merge into pkg/tool

cmdprof functionality seems to fit well into pkg/tool.

4 files changed, 67 insertions, 16 deletions

diff --git a/pkg/tool/cmdprof.go b/pkg/tool/cmdprof.go
new file mode 100644
index 000000000..196a0cd24
--- /dev/null
+++ b/pkg/tool/cmdprof.go
@@ -0,0 +1,44 @@
+// Copyright 2020 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package tool
+
+import (
+	"os"
+	"runtime"
+	"runtime/pprof"
+)
+
+// installProfiling simplifies cpu/memory profiling for command line tools.
+func installProfiling(cpuprof, memprof string) func() {
+	res := func() {}
+	if cpuprof != "" {
+		f, err := os.Create(cpuprof)
+		if err != nil {
+			Failf("failed to create cpuprofile file: %v", err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			Failf("failed to start cpu profile: %v", err)
+		}
+		res = func() {
+			pprof.StopCPUProfile()
+			f.Close()
+		}
+	}
+	if memprof != "" {
+		prev := res
+		res = func() {
+			prev()
+			f, err := os.Create(memprof)
+			if err != nil {
+				Failf("failed to create memprofile file: %v", err)
+			}
+			defer f.Close()
+			runtime.GC()
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				Failf("failed to write mem profile: %v", err)
+			}
+		}
+	}
+	return res
+}
diff --git a/pkg/tool/flags.go b/pkg/tool/flags.go
index 197494823..0763cb502 100644
--- a/pkg/tool/flags.go
+++ b/pkg/tool/flags.go
@@ -8,39 +8,31 @@ import (
 	"encoding/hex"
 	"flag"
 	"fmt"
-	"os"
 	"strings"
 
 	"github.com/google/syzkaller/pkg/log"
 )
 
-// ParseFlags parses command line flags with flag.Parse and then applies optional flags created with OptionalFlags.
-// This is intended for programmatic use only when we invoke older versions of binaries with new unsupported flags.
-func ParseFlags() {
-	if err := parseFlags(flag.CommandLine, os.Args[1:]); err != nil {
-		Fail(err)
-	}
-}
-
 type Flag struct {
 	Name  string
 	Value string
 }
 
 // OptionalFlags produces command line flag value that encapsulates the given flags as optional.
-// Use ParseFlags to support optional flags in the binary.
+// This is intended for programmatic use only when we invoke older versions of binaries with new unsupported flags.
+// Use tool.Init to support optional flags in the binary.
 // The format keeps flags reasonably readable ("-optional=foo=bar:baz=123"), not subject to accidental splitting
 // into multiple arguments due to spaces and supports bool/non-bool flags.
-func OptionalFlags(flags ...Flag) string {
+func OptionalFlags(flags []Flag) string {
 	return fmt.Sprintf("-%v=%v", optionalFlag, serializeFlags(flags))
 }
 
-func parseFlags(set *flag.FlagSet, args []string) error {
-	optional := set.String(optionalFlag, "", "optional flags for programmatic use only")
+func ParseFlags(set *flag.FlagSet, args []string) error {
+	flagOptional := set.String(optionalFlag, "", "optional flags for programmatic use only")
 	if err := set.Parse(args); err != nil {
 		return err
 	}
-	flags, err := deserializeFlags(*optional)
+	flags, err := deserializeFlags(*flagOptional)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/tool/flags_test.go b/pkg/tool/flags_test.go
index 6d52a4780..d818c6582 100644
--- a/pkg/tool/flags_test.go
+++ b/pkg/tool/flags_test.go
@@ -27,7 +27,7 @@ func TestParseFlags(t *testing.T) {
 		{"", &Values{false, 1, "baz"}},
 		{"-foo -bar=2", &Values{true, 2, "baz"}},
 		{"-foo -bar=2 -qux", nil},
-		{"-foo -bar=2 " + OptionalFlags(Flag{"qux", ""}), &Values{true, 2, "baz"}},
+		{"-foo -bar=2 " + OptionalFlags([]Flag{{"qux", ""}}), &Values{true, 2, "baz"}},
 	}
 	for i, test := range tests {
 		t.Run(fmt.Sprint(i), func(t *testing.T) {
@@ -41,7 +41,7 @@ func TestParseFlags(t *testing.T) {
 			if args[0] == "" {
 				args = args[1:]
 			}
-			err := parseFlags(flags, args)
+			err := ParseFlags(flags, args)
 			if test.vals == nil {
 				if err == nil {
 					t.Fatalf("parsing did not fail")
diff --git a/pkg/tool/tool.go b/pkg/tool/tool.go
index bb1f436cc..690bfa4b8 100644
--- a/pkg/tool/tool.go
+++ b/pkg/tool/tool.go
@@ -5,10 +5,25 @@
 package tool
 
 import (
+	"flag"
 	"fmt"
 	"os"
 )
 
+// Init handles common tasks for command line tools:
+//  - invokes flag.Parse
+//  - adds support for optional flags (see OptionalFlags)
+//  - adds support for cpu/mem profiling (-cpuprofile/memprofile flags)
+// Use as defer tool.Init()().
+func Init() func() {
+	flagCPUProfile := flag.String("cpuprofile", "", "write CPU profile to this file")
+	flagMEMProfile := flag.String("memprofile", "", "write memory profile to this file")
+	if err := ParseFlags(flag.CommandLine, os.Args[1:]); err != nil {
+		Fail(err)
+	}
+	return installProfiling(*flagCPUProfile, *flagMEMProfile)
+}
+
 func Failf(msg string, args ...interface{}) {
 	fmt.Fprintf(os.Stderr, msg+"\n", args...)
 	os.Exit(1)