Merge branch 'ortz-feature/timeout-flag'

emitter.go

@@ -32,13 +32,7 @@ func Start(stop chan int) {
 	for {
 		select {
 		case <-stop:
-			pluginMu.Lock()
-			for _, p := range Plugins.All {
-				if cp, ok := p.(io.Closer); ok {
-					cp.Close()
-				}
-			}
-			pluginMu.Unlock()
+			finalize()
 			return
 		case <-time.After(100 * time.Millisecond):
 		}

gor.go

@@ -78,17 +78,31 @@ func main() {
 	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
 	go func() {
 		<-c
-
-		for _, p := range Plugins.All {
-			if cp, ok := p.(io.Closer); ok {
-				cp.Close()
-			}
-		}
-
+		finalize()
 		os.Exit(1)
 	}()
 
-	Start(nil)
+	if Settings.exitAfter > 0 {
+		log.Println("Running gor for a duration of", Settings.exitAfter)
+		closeCh := make(chan int)
+
+		time.AfterFunc(Settings.exitAfter, func() {
+			log.Println("Stopping gor after", Settings.exitAfter)
+			close(closeCh)
+		})
+
+		Start(closeCh)
+	} else {
+		Start(nil)
+	}
+}
+
+func finalize() {
+	for _, p := range Plugins.All {
+		if cp, ok := p.(io.Closer); ok {
+			cp.Close()
+		}
+	}
 }
 
 func profileCPU(cpuprofile string) {
@@ -118,4 +132,4 @@ func profileMEM(memprofile string) {
 			f.Close()
 		})
 	}
-}
+}

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.")