diff --git a/emitter.go b/emitter.go index 9575484..1fc2a83 100644 --- a/emitter.go +++ b/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): } diff --git a/gor.go b/gor.go index 8788a74..085eb98 100644 --- a/gor.go +++ b/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() }) } -} +} \ No newline at end of file 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.")