@@ 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 {