Refactor to keep printing decoupled from stress run (#26)

cmd/stress.go

@@ -129,7 +129,7 @@ var stressCmd = &cobra.Command{
 				//info about the request
 				fmt.Printf("----Target %d: %s %s\n", idx+1, target.Method, target.URL)
 				reqStats := pewpew.CreateRequestsStats(targetRequestStats[idx])
-				fmt.Println(pewpew.CreateTextSummary(reqStats))
+				fmt.Println(pewpew.CreateTextStressSummary(reqStats))
 			}
 		}
 
@@ -144,7 +144,7 @@ var stressCmd = &cobra.Command{
 			fmt.Println("----Global----")
 		}
 		reqStats := pewpew.CreateRequestsStats(globalStats)
-		fmt.Println(pewpew.CreateTextSummary(reqStats))
+		fmt.Println(pewpew.CreateTextStressSummary(reqStats))
 
 		//write out json
 		if viper.GetString("ResultFilenameJSON") != "" {

lib/printer.go

@@ -6,12 +6,21 @@ import (
 	"io/ioutil"
 	"net/http"
 	"sort"
+	"sync"
 
 	color "github.com/fatih/color"
 )
 
-//CreateTextSummary creates a human friendly summary of entire stress test
-func CreateTextSummary(reqStatSummary RequestStatSummary) string {
+type printer struct {
+	//writeLock prevents concurrent messages from being interlaced
+	writeLock sync.Mutex
+
+	//output is where the printer writes to
+	output io.Writer
+}
+
+//CreateTextStressSummary creates a human friendly summary of entire stress test
+func CreateTextStressSummary(reqStatSummary RequestStatSummary) string {
 	summary := "\n"
 
 	summary += "Timing\n"
@@ -54,36 +63,40 @@ func CreateTextSummary(reqStatSummary RequestStatSummary) string {
 }
 
 //print colored single line stats per RequestStat
-func printStat(stat RequestStat, w io.Writer) {
+func (p *printer) printStat(stat RequestStat) {
+	p.writeLock.Lock()
+	defer p.writeLock.Unlock()
+
 	if stat.Error != nil {
 		color.Set(color.FgRed)
-		fmt.Fprintln(w, "Failed to make request: "+stat.Error.Error())
+		fmt.Fprintln(p.output, "Failed to make request: "+stat.Error.Error())
 		color.Unset()
+		return
+	}
+
+	if stat.StatusCode >= 100 && stat.StatusCode < 200 {
+		color.Set(color.FgBlue)
+	} else if stat.StatusCode >= 200 && stat.StatusCode < 300 {
+		color.Set(color.FgGreen)
+	} else if stat.StatusCode >= 300 && stat.StatusCode < 400 {
+		color.Set(color.FgCyan)
+	} else if stat.StatusCode >= 400 && stat.StatusCode < 500 {
+		color.Set(color.FgMagenta)
 	} else {
-		if stat.StatusCode >= 100 && stat.StatusCode < 200 {
-			color.Set(color.FgBlue)
-		} else if stat.StatusCode >= 200 && stat.StatusCode < 300 {
-			color.Set(color.FgGreen)
-		} else if stat.StatusCode >= 300 && stat.StatusCode < 400 {
-			color.Set(color.FgCyan)
-		} else if stat.StatusCode >= 400 && stat.StatusCode < 500 {
-			color.Set(color.FgMagenta)
-		} else {
-			color.Set(color.FgRed)
-		}
-		fmt.Fprintf(w, "%s %d\t%d bytes\t%d ms\t-> %s %s\n",
-			stat.Proto,
-			stat.StatusCode,
-			stat.DataTransferred,
-			stat.Duration.Nanoseconds()/1000000,
-			stat.Method,
-			stat.URL)
-		color.Unset()
+		color.Set(color.FgRed)
 	}
+	fmt.Fprintf(p.output, "%s %d\t%d bytes\t%d ms\t-> %s %s\n",
+		stat.Proto,
+		stat.StatusCode,
+		stat.DataTransferred,
+		stat.Duration.Nanoseconds()/1000000,
+		stat.Method,
+		stat.URL)
+	color.Unset()
 }
 
 //print tons of info about the request, response and response body
-func printVerbose(req *http.Request, response *http.Response, w io.Writer) {
+func (p *printer) printVerbose(req *http.Request, response *http.Response) {
 	if req == nil {
 		return
 	}
@@ -105,5 +118,14 @@ func printVerbose(req *http.Request, response *http.Response, w io.Writer) {
 		requestInfo = requestInfo + fmt.Sprintf("Body:\n%s\n\n", body)
 		response.Body.Close()
 	}
-	fmt.Fprintln(w, requestInfo)
+	p.writeLock.Lock()
+	fmt.Fprintln(p.output, requestInfo)
+	p.writeLock.Unlock()
+}
+
+//writeString is a generic output string printer
+func (p *printer) writeString(s string) {
+	p.writeLock.Lock()
+	fmt.Fprint(p.output, s)
+	p.writeLock.Unlock()
 }

lib/printer_test.go

@@ -29,7 +29,7 @@ func TestCreateTextSummary(t *testing.T) {
 	}
 	for _, c := range cases {
 		//could check for the exact string, but that's super tedious and brittle
-		_ = CreateTextSummary(c.s)
+		_ = CreateTextStressSummary(c.s)
 	}
 }
 
@@ -47,8 +47,9 @@ func TestPrintStat(t *testing.T) {
 		//error case
 		{RequestStat{Error: errors.New("this is an error")}},
 	}
+	p := printer{output: ioutil.Discard}
 	for _, c := range cases {
-		printStat(c.r, ioutil.Discard)
+		p.printStat(c.r)
 	}
 }
 
@@ -62,7 +63,8 @@ func TestPrintVerbose(t *testing.T) {
 		{&http.Request{}, nil},
 		{&http.Request{}, &http.Response{Body: http.NoBody}},
 	}
+	p := printer{output: ioutil.Discard}
 	for _, c := range cases {
-		printVerbose(c.req, c.resp, ioutil.Discard)
+		p.printVerbose(c.req, c.resp)
 	}
 }

lib/stress.go

@@ -5,13 +5,8 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"sync"
 )
 
-//TODO move to other file
-//so concurrent workers don't interlace messages
-var writeLock sync.Mutex
-
 //TODO move to other file
 type workerDone struct{}
 
@@ -77,6 +72,9 @@ func RunStress(s StressConfig, w io.Writer) ([][]RequestStat, error) {
 	}
 	targetCount := len(s.Targets)
 
+	//setup printer
+	p := printer{output: w}
+
 	//setup the queue of requests, one queue per target
 	requestQueues := make([](chan http.Request), targetCount)
 	for idx, target := range s.Targets {
@@ -101,9 +99,7 @@ func RunStress(s StressConfig, w io.Writer) ([][]RequestStat, error) {
 	targetStats := make(chan []RequestStat)
 	for idx, target := range s.Targets {
 		go func(target Target, requestQueue chan http.Request, targetStats chan []RequestStat) {
-			writeLock.Lock()
-			fmt.Fprintf(w, "- Running %d tests at %s, %d at a time\n", s.Count, target.URL, s.Concurrency)
-			writeLock.Unlock()
+			p.writeString(fmt.Sprintf("- Running %d tests at %s, %d at a time\n", s.Count, target.URL, s.Concurrency))
 
 			workerDoneChan := make(chan workerDone)   //workers use this to indicate they are done
 			requestStatChan := make(chan RequestStat) //workers communicate each requests' info
@@ -124,12 +120,10 @@ func RunStress(s StressConfig, w io.Writer) ([][]RequestStat, error) {
 
 							response, stat := runRequest(req, client)
 							if !s.Quiet {
-								writeLock.Lock()
-								printStat(stat, w)
+								p.printStat(stat)
 								if s.Verbose {
-									printVerbose(&req, response, w)
+									p.printVerbose(&req, response)
 								}
-								writeLock.Unlock()
 							}
 
 							requestStatChan <- stat