Merge pull request #18 from PouuleT/asyncSignal

Async signal

Author: PouuleT

http_stats.go

@@ -89,7 +89,7 @@ func (s *HTTPStats) Render() {
 		strconv.Itoa(s.nbOfRequests),
 		s.minDuration.String(),
 		s.maxDuration.String(),
-		(s.totalDuration / time.Duration(s.nbOfRequests)).String(),
+		getAvgDuration(s.totalDuration, s.nbOfRequests),
 		s.execDuration.String(),
 		fmt.Sprintf("%s/s", humanize.Bytes(uint64(float64(time.Duration(s.totalSize))/float64(s.execDuration)*math.Pow10(9)))),
 		fmt.Sprintf("%s", humanize.Bytes(uint64(s.totalSize))),
@@ -112,19 +112,19 @@ func (s *HTTPStats) Render() {
 	timeTable.SetHeader([]string{"Step", "Average duration"})
 	timeTable.SetAlignment(tablewriter.ALIGN_CENTER)
 	timeTable.Append([]string{
-		"DNSLookup", (s.responseTimeline.DNSLookup / time.Duration(s.successRequests)).String(),
+		"DNSLookup", getAvgDuration(s.responseTimeline.DNSLookup, s.successRequests),
 	})
 	timeTable.Append([]string{
-		"TCPConnection", (s.responseTimeline.TCPConnection / time.Duration(s.successRequests)).String(),
+		"TCPConnection", getAvgDuration(s.responseTimeline.TCPConnection, s.successRequests),
 	})
 	timeTable.Append([]string{
-		"EstablishingConnection", (s.responseTimeline.EstablishingConnection / time.Duration(s.successRequests)).String(),
+		"EstablishingConnection", getAvgDuration(s.responseTimeline.EstablishingConnection, s.successRequests),
 	})
 	timeTable.Append([]string{
-		"ServerProcessing", (s.responseTimeline.ServerProcessing / time.Duration(s.successRequests)).String(),
+		"ServerProcessing", getAvgDuration(s.responseTimeline.ServerProcessing, s.successRequests),
 	})
 	timeTable.Append([]string{
-		"ContentTransfer", (s.responseTimeline.ContentTransfer / time.Duration(s.successRequests)).String(),
+		"ContentTransfer", getAvgDuration(s.responseTimeline.ContentTransfer, s.successRequests),
 	})
 
 	fmt.Printf("\nRequest details :\n")

traffic_generator.go

@@ -5,7 +5,9 @@ import (
 	"log"
 	"math"
 	"os"
+	"os/signal"
 	"sync"
+	"syscall"
 	"time"
 )
 
@@ -32,6 +34,8 @@ var statsMap = map[string]func() Stats{
 	"dns":  newDNSStats,
 }
 
+var exitChan = make(chan struct{})
+
 // NewTrafficGenerator will return a new TrafficGenerator object
 func NewTrafficGenerator(trafficType string) (*TrafficGenerator, error) {
 	tFunc, ok := trafficMap[trafficType]
@@ -50,11 +54,53 @@ func NewTrafficGenerator(trafficType string) (*TrafficGenerator, error) {
 
 // Generate generates traffic
 func (trafficGen *TrafficGenerator) Generate() {
+	// Create a channel that will listen to SIGINT / SIGTERM
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGINT)
+	signal.Notify(c, syscall.SIGTERM)
+
 	for i := 1; i <= nbOfClients; i++ {
 		trafficGen.wg.Add(1)
+		// Launch the workers in a go routine
 		go trafficGen.NewWorker(i).work()
 	}
-	trafficGen.wg.Wait()
+
+	// Done channel to stop the loop when all the workers are done
+	var done = make(chan struct{})
+	go func() {
+		// Wait for the workerss
+		trafficGen.wg.Wait()
+		// All the workers are done
+		done <- struct{}{}
+	}()
+
+	// Wait for the workers to end
+	// or a signal in the loop
+	var forceShutdown bool
+	for {
+		select {
+		case <-done:
+			// All the workers are done, we quit
+			return
+		case sig := <-c:
+			// We listen for signals
+			switch sig {
+			case syscall.SIGINT, syscall.SIGKILL:
+				// If it's the second time we get a signal, quit
+				if forceShutdown {
+					os.Exit(1)
+				}
+
+				// Notify all the workers that they need to stop
+				for i := 1; i <= nbOfClients; i++ {
+					exitChan <- struct{}{}
+				}
+
+				// Next time we get a signal, need to quit directly
+				forceShutdown = true
+			}
+		}
+	}
 }
 
 // NewWorker creates a new worker for traffic generation
@@ -71,16 +117,40 @@ func (trafficGen *TrafficGenerator) DisplayStats() {
 }
 
 func (w *Worker) work() {
+	var exit bool
 	defer w.trafficGen.wg.Done()
 	start := time.Now()
 
+	var done = make(chan struct{})
+	// When the work is done, notify the watching go routine
+	defer func() { done <- struct{}{} }()
+
+	// Create the watching go routine that will watch if we need to quit early
+	go func() {
+		for {
+			select {
+			// If we need to exit
+			case <-exitChan:
+				exit = true
+			// If everything is done
+			case <-done:
+				return
+			}
+		}
+	}()
+
 	workerFmt := fmt.Sprintf("worker#%%0%dd", getPadding(nbOfClients))
 	counterFmt := fmt.Sprintf(" - %%0%dd/%%d ", getPadding(nbOfRequests))
 
 	prefix := fmt.Sprintf(workerFmt, w.id)
 	logger := log.New(os.Stdout, prefix, 0)
+
 	// Repeat nbOfRequests requests
 	for i := 1; i <= nbOfRequests; i++ {
+		// If we got an exit signal, quit
+		if exit {
+			return
+		}
 		logger.SetPrefix(prefix + fmt.Sprintf(counterFmt, i, nbOfRequests))
 		// Find an URL
 		url := findRandomURL()
@@ -101,3 +171,10 @@ func getPadding(nb int) int {
 	// Get the padding size : floor(log10(nb)) + 1
 	return int(math.Log10(float64(nb))) + 1
 }
+
+func getAvgDuration(total time.Duration, number int) string {
+	if number == 0 {
+		return "NaN"
+	}
+	return (total / time.Duration(number)).String()
+}