From 6620d256264802aa4d67d26ad92437a758be4ff8 Mon Sep 17 00:00:00 2001
From: Or Tzabary <or@tzabary.com>
Date: Fri, 8 Jul 2016 18:57:42 +0300
Subject: [PATCH] Change flag name, use of DurationVar instead of int

---
 gor.go      | 13 ++++++-------
 settings.go |  8 +++++---
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/gor.go b/gor.go
index c0f06c7..70f3fa7 100644
--- a/gor.go
+++ b/gor.go
@@ -22,7 +22,6 @@ var (
 	mode       string
 	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
 	memprofile = flag.String("memprofile", "", "write memory profile to this file")
-	timeout    = flag.Int("timeout", 0, "timeout to stop gor, value in seconds")
 )
 
 func loggingMiddleware(next http.Handler) http.Handler {
@@ -89,10 +88,10 @@ func main() {
 		os.Exit(1)
 	}()
 
-	if *timeout >= 1 {
-		log.Println("Running gor with timeout of", *timeout, "seconds")
+	if Settings.exitAfter >= 1 {
+		log.Println("Running gor for a duration of", Settings.exitAfter)
 		stop := make(chan int)
-		timeoutGor(stop, *timeout)
+		stopAfter(stop, Settings.exitAfter)
 
 		Start(stop)
 	} else {
@@ -129,9 +128,9 @@ func profileMEM(memprofile string) {
 	}
 }
 
-func timeoutGor(stop chan int, seconds int) {
-	time.AfterFunc(time.Duration(seconds)*time.Second, func() {
-		log.Println("Stopping gor after", seconds, "seconds")
+func stopAfter(stop chan int, exitAfter time.Duration) {
+	time.AfterFunc(exitAfter, func() {
+		log.Println("Stopping gor, duration of", exitAfter, "reached")
 		close(stop)
 	})
 }
diff --git a/settings.go b/settings.go
index dd0675c..1b6104e 100644
--- a/settings.go
+++ b/settings.go
@@ -25,9 +25,10 @@ func (h *MultiOption) Set(value string) error {
 
 // AppSettings is the struct of main configuration
 type AppSettings struct {
-	verbose bool
-	debug   bool
-	stats   bool
+	verbose   bool
+	debug     bool
+	stats     bool
+	exitAfter time.Duration
 
 	splitOutput bool
 
@@ -74,6 +75,7 @@ func init() {
 	flag.BoolVar(&Settings.verbose, "verbose", false, "Turn on more verbose output")
 	flag.BoolVar(&Settings.debug, "debug", false, "Turn on debug output, shows all intercepted traffic. Works only when with `verbose` flag")
 	flag.BoolVar(&Settings.stats, "stats", false, "Turn on queue stats output")
+	flag.DurationVar(&Settings.exitAfter, "exit-after", 0, "exit after specified duration")
 
 	flag.BoolVar(&Settings.splitOutput, "split-output", false, "By default each output gets same traffic. If set to `true` it splits traffic equally among all outputs.")