From dae6a8915de74aaa08f6d697d20037cdb353f804 Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Fri, 8 Jul 2016 19:27:20 +0300
Subject: [PATCH] Refactor exit-after

---
 emitter.go |  8 +-------
 gor.go     | 37 ++++++++++++++++++-------------------
 2 files changed, 19 insertions(+), 26 deletions(-)

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 70f3fa7..085eb98 100644
--- a/gor.go
+++ b/gor.go
@@ -78,27 +78,33 @@ 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)
 	}()
 
-	if Settings.exitAfter >= 1 {
+	if Settings.exitAfter > 0 {
 		log.Println("Running gor for a duration of", Settings.exitAfter)
-		stop := make(chan int)
-		stopAfter(stop, Settings.exitAfter)
+		closeCh := make(chan int)
 
-		Start(stop)
+		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) {
 	if cpuprofile != "" {
 		f, err := os.Create(cpuprofile)
@@ -126,11 +132,4 @@ func profileMEM(memprofile string) {
 			f.Close()
 		})
 	}
-}
-
-func stopAfter(stop chan int, exitAfter time.Duration) {
-	time.AfterFunc(exitAfter, func() {
-		log.Println("Stopping gor, duration of", exitAfter, "reached")
-		close(stop)
-	})
-}
+}
\ No newline at end of file