1 files changed, 25 insertions(+), 2 deletions(-)

M main.go
M main.go +25 -2
@@ 8,6 8,7 @@ import (
 	"math/rand"
 	"net/http"
 	"os"
+	"os/signal"
 	"sync"
 	"time"
 

          
@@ 34,6 35,7 @@ func worker(
 	duration time.Duration,
 	output chan<- Result,
 	waitGroup *sync.WaitGroup,
+	terminate <-chan struct{},
 ) {
 	defer waitGroup.Done()
 

          
@@ 42,10 44,18 @@ func worker(
 		logger.Println("Error connecting to", seleniumUrl, err)
 		return
 	}
+	defer driver.Quit()
 
 	var startTime = time.Now()
 
+loop:
 	for time.Since(startTime) < duration {
+		select {
+		case <-terminate:
+			logger.Println("worker exit")
+			break loop
+		default:
+		}
 		var n = time.Now()
 		err = driver.Get(url)
 		var d = time.Since(n)

          
@@ 57,7 67,6 @@ func worker(
 			output <- Result{Timestamp: n, Duration: d}
 		}
 	}
-	driver.Quit()
 }
 
 //

          
@@ 149,13 158,15 @@ func main() {
 
 	// Buffer results to avoid blocking workers
 	var output = make(chan Result, 10)
+	// control termination of workers goroutines if we press Ctrl+C
+	var terminate = make(chan struct{}, workers)
 	var wg sync.WaitGroup
 
 	logger.Println("Using", workers, "workers for", duration, "seconds")
 
 	wg.Add(workers)
 	for i := 0; i < workers; i++ {
-		go worker(capabilities, url, remoteUrl, time.Duration(duration)*time.Second, output, &wg)
+		go worker(capabilities, url, remoteUrl, time.Duration(duration)*time.Second, output, &wg, terminate)
 	}
 
 	// Wait for all worker to finish, then close output to end the program

          
@@ 164,6 175,18 @@ func main() {
 		close(output)
 	}()
 
+	// Wait for an interupt signal, terminate workers if needed
+	go func() {
+		var c = make(chan os.Signal, 1)
+		signal.Notify(c, os.Interrupt)
+
+		s := <-c
+		logger.Println("Got signal", s, "terminating workers")
+		for i := 0; i < workers; i++ {
+			terminate <- struct{}{}
+		}
+	}()
+
 	var writer = csv.NewWriter(os.Stdout)
 
 	for o := range output {