Refactor exit-after

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,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)
-	})
-}
+}