# HG changeset patch
# User Henry Precheur <henry@precheur.org>
# Date 1476498350 25200
#      Fri Oct 14 19:25:50 2016 -0700
# Node ID 51641a0e2318cfba42f5d1f1c7fe6cc496a9bf5d
# Parent  3979a82a2f66d4d0a672a5fe2ea3426bf3e1e0b9
Handle ctrl+c

diff --git a/main.go b/main.go
--- a/main.go
+++ b/main.go
@@ -8,6 +8,7 @@
 	"math/rand"
 	"net/http"
 	"os"
+	"os/signal"
 	"sync"
 	"time"
 
@@ -34,6 +35,7 @@
 	duration time.Duration,
 	output chan<- Result,
 	waitGroup *sync.WaitGroup,
+	terminate <-chan struct{},
 ) {
 	defer waitGroup.Done()
 
@@ -42,10 +44,18 @@
 		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 @@
 			output <- Result{Timestamp: n, Duration: d}
 		}
 	}
-	driver.Quit()
 }
 
 //
@@ -149,13 +158,15 @@
 
 	// 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 @@
 		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 {