pkg/cmdprof: merge into pkg/tool

cmdprof functionality seems to fit well into pkg/tool.

1 files changed, 15 insertions, 0 deletions

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)