Added profiling to CLI

asticode · Nov 26, 2017 · 207bfb4 · 207bfb4
1 parent 5d9f513
commit 207bfb4
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/astits/main.go b/astits/main.go
@@ -18,17 +18,24 @@ import (
 	"github.com/asticode/go-astitools/flag"
 	"github.com/asticode/go-astits"
 	"github.com/pkg/errors"
+	"github.com/pkg/profile"
 )
 
 // Flags
 var (
-	ctx, cancel = context.WithCancel(context.Background())
-	dataTypes   = astiflag.NewStringsMap()
-	format      = flag.String("f", "", "the format")
-	inputPath   = flag.String("i", "", "the input path")
+	ctx, cancel    = context.WithCancel(context.Background())
+	dataTypes      = astiflag.NewStringsMap()
+	format         = flag.String("f", "", "the format")
+	inputPath      = flag.String("i", "", "the input path")
+	profileEnabled = flag.Bool("p", false, "if yes, profiling is enabled")
 )
 
 func main() {
+	// Start profiling
+	if *profileEnabled {
+		defer profile.Start(profile.MemProfile, profile.CPUProfile, profile.ProfilePath(".")).Stop()
+	}
+
 	// Init
 	flag.Var(dataTypes, "d", "the datatypes whitelist")
 	var s = astiflag.Subcommand()