Add --concurrent option and support concurrent requests

Author: maddie

defs/bytes_counter.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"sync"
 	"time"
 )
 
@@ -16,20 +17,32 @@ type BytesCounter struct {
 	payload []byte
 	reader  io.ReadSeeker
 	mebi    bool
+
+	lock *sync.Mutex
+}
+
+func NewCounter() *BytesCounter {
+	return &BytesCounter{
+		lock: &sync.Mutex{},
+	}
 }
 
 // Write implements io.Writer
 func (c *BytesCounter) Write(p []byte) (int, error) {
 	n := len(p)
+	c.lock.Lock()
 	c.total += n
+	c.lock.Unlock()
 
 	return n, nil
 }
 
 // Read implements io.Reader
 func (c *BytesCounter) Read(p []byte) (int, error) {
 	n, err := c.reader.Read(p)
+	c.lock.Lock()
 	c.total += n
+	c.lock.Unlock()
 
 	return n, err
 }
@@ -39,8 +52,8 @@ func (c *BytesCounter) SetMebi(mebi bool) {
 	c.mebi = mebi
 }
 
-// Average returns the average bytes/second
-func (c *BytesCounter) Average() float64 {
+// AvgBytes returns the average bytes/second
+func (c *BytesCounter) AvgBytes() float64 {
 	return float64(c.total) / time.Now().Sub(c.start).Seconds()
 }
 
@@ -49,11 +62,11 @@ func (c *BytesCounter) AvgMbps() float64 {
 	if c.mebi {
 		base = 131072
 	}
-	return c.Average() / base
+	return c.AvgBytes() / base
 }
 
 func (c *BytesCounter) AvgHumanize() string {
-	val := c.Average()
+	val := c.AvgBytes()
 
 	var base float64 = 1000
 	if c.mebi {

defs/options.go

@@ -8,6 +8,7 @@ const (
 	OptionIPv6Alt         = "6"
 	OptionNoDownload      = "no-download"
 	OptionNoUpload        = "no-upload"
+	OptionConcurrent      = "concurrent"
 	OptionBytes           = "bytes"
 	OptionMebiBytes       = "mebibytes"
 	OptionDistance        = "distance"

defs/server.go

@@ -177,13 +177,13 @@ func (s *Server) PingAndJitter(count int) (float64, float64, error) {
 }
 
 // Download performs the actual download test
-func (s *Server) Download(silent bool, useBytes, useMebi bool) (float64, int, error) {
+func (s *Server) Download(silent bool, useBytes, useMebi bool, requests int) (float64, int, error) {
 	t := time.Now()
 	defer func() {
 		s.TLog.Logf("Download took %s", time.Now().Sub(t).String())
 	}()
 
-	counter := &BytesCounter{}
+	counter := NewCounter()
 	counter.SetMebi(useMebi)
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -207,7 +207,7 @@ func (s *Server) Download(silent bool, useBytes, useMebi bool) (float64, int, er
 	req.Header.Set("User-Agent", UserAgent)
 	req.Header.Set("Accept-Encoding", "identity")
 
-	downloadDone := make(chan struct{})
+	downloadDone := make(chan struct{}, requests)
 
 	doDownload := func() {
 		resp, err := http.DefaultClient.Do(req)
@@ -249,7 +249,10 @@ func (s *Server) Download(silent bool, useBytes, useMebi bool) (float64, int, er
 		}()
 	}
 
-	go doDownload()
+	for i := 0; i < requests; i++ {
+		go doDownload()
+		time.Sleep(200 * time.Millisecond)
+	}
 	timeout := time.After(15 * time.Second)
 Loop:
 	for {
@@ -266,13 +269,13 @@ Loop:
 }
 
 // Upload performs the actual upload test
-func (s *Server) Upload(noPrealloc, silent, useBytes, useMebi bool) (float64, int, error) {
+func (s *Server) Upload(noPrealloc, silent, useBytes, useMebi bool, requests int) (float64, int, error) {
 	t := time.Now()
 	defer func() {
 		s.TLog.Logf("Upload took %s", time.Now().Sub(t).String())
 	}()
 
-	counter := &BytesCounter{}
+	counter := NewCounter()
 	counter.SetMebi(useMebi)
 
 	if noPrealloc {
@@ -299,7 +302,7 @@ func (s *Server) Upload(noPrealloc, silent, useBytes, useMebi bool) (float64, in
 	req.Header.Set("User-Agent", UserAgent)
 	req.Header.Set("Accept-Encoding", "identity")
 
-	uploadDone := make(chan struct{})
+	uploadDone := make(chan struct{}, requests)
 
 	doUpload := func() {
 		resp, err := http.DefaultClient.Do(req)
@@ -339,7 +342,10 @@ func (s *Server) Upload(noPrealloc, silent, useBytes, useMebi bool) (float64, in
 		}()
 	}
 
-	go doUpload()
+	for i := 0; i < requests; i++ {
+		go doUpload()
+		time.Sleep(200 * time.Millisecond)
+	}
 	timeout := time.After(15 * time.Second)
 Loop:
 	for {

main.go

@@ -55,6 +55,11 @@ func main() {
 				Name:  defs.OptionNoUpload,
 				Usage: "Do not perform upload test",
 			},
+			&cli.IntFlag{
+				Name:  defs.OptionConcurrent,
+				Usage: "Concurrent HTTP requests being made",
+				Value: 3,
+			},
 			&cli.BoolFlag{
 				Name: defs.OptionBytes,
 				Usage: "Display values in bytes instead of bits. Does not affect\n" +

speedtest/helper.go

@@ -78,7 +78,7 @@ func doSpeedTest(c *cli.Context, servers []defs.Server, telemetryServer defs.Tel
 			if c.Bool(defs.OptionNoDownload) {
 				log.Info("Download test is disabled")
 			} else {
-				download, br, err := currentServer.Download(silent, c.Bool(defs.OptionBytes), c.Bool(defs.OptionMebiBytes))
+				download, br, err := currentServer.Download(silent, c.Bool(defs.OptionBytes), c.Bool(defs.OptionMebiBytes), c.Int(defs.OptionConcurrent))
 				if err != nil {
 					log.Errorf("Failed to get download speed: %s", err)
 					return err
@@ -93,7 +93,7 @@ func doSpeedTest(c *cli.Context, servers []defs.Server, telemetryServer defs.Tel
 			if c.Bool(defs.OptionNoUpload) {
 				log.Info("Upload test is disabled")
 			} else {
-				upload, bw, err := currentServer.Upload(c.Bool(defs.OptionNoPreAllocate), silent, c.Bool(defs.OptionBytes), c.Bool(defs.OptionMebiBytes))
+				upload, bw, err := currentServer.Upload(c.Bool(defs.OptionNoPreAllocate), silent, c.Bool(defs.OptionBytes), c.Bool(defs.OptionMebiBytes), c.Int(defs.OptionConcurrent))
 				if err != nil {
 					log.Errorf("Failed to get upload speed: %s", err)
 					return err

speedtest/speedtest.go

@@ -123,6 +123,11 @@ func SpeedTest(c *cli.Context) error {
 		}
 	}
 
+	if req := c.Int(defs.OptionConcurrent); req <= 0 {
+		log.Errorf("Concurrent requests cannot be lower than 1: %d is given", req)
+		return errors.New("invalid concurrent requests setting")
+	}
+
 	// HTTP requests timeout
 	http.DefaultClient.Timeout = time.Duration(c.Int(defs.OptionTimeout)) * time.Second