Configurable download chunks, upload size and test duration

Author: maddie

defs/bytes_counter.go

@@ -12,12 +12,13 @@ import (
 
 // BytesCounter implements io.Reader and io.Writer interface, for counting bytes being read/written in HTTP requests
 type BytesCounter struct {
-	start   time.Time
-	pos     int
-	total   int
-	payload []byte
-	reader  io.ReadSeeker
-	mebi    bool
+	start      time.Time
+	pos        int
+	total      int
+	payload    []byte
+	reader     io.ReadSeeker
+	mebi       bool
+	uploadSize int
 
 	lock *sync.Mutex
 }
@@ -44,7 +45,7 @@ func (c *BytesCounter) Read(p []byte) (int, error) {
 	c.lock.Lock()
 	c.total += n
 	c.pos += n
-	if c.pos == uploadSize {
+	if c.pos == c.uploadSize {
 		c.resetReader()
 	}
 	c.lock.Unlock()
@@ -57,6 +58,11 @@ func (c *BytesCounter) SetMebi(mebi bool) {
 	c.mebi = mebi
 }
 
+// SetUploadSize sets the size of payload being uploaded
+func (c *BytesCounter) SetUploadSize(uploadSize int) {
+	c.uploadSize = uploadSize * 1024
+}
+
 // AvgBytes returns the average bytes/second
 func (c *BytesCounter) AvgBytes() float64 {
 	return float64(c.total) / time.Now().Sub(c.start).Seconds()
@@ -94,7 +100,7 @@ func (c *BytesCounter) AvgHumanize() string {
 // GenerateBlob generates a random byte array of `uploadSize` in the `payload` field, and sets the `reader` field to
 // read from it
 func (c *BytesCounter) GenerateBlob() {
-	c.payload = getRandomData(uploadSize)
+	c.payload = getRandomData(c.uploadSize)
 	c.reader = bytes.NewReader(c.payload)
 }
 

defs/defs.go

@@ -1,12 +1,5 @@
 package defs
 
-const (
-	// chunks to download in download test
-	downloadChunks = 100
-	// payload size per upload request
-	uploadSize = 1024 * 1024
-)
-
 var (
 	// values to be filled in by build script
 	BuildDate   string

defs/options.go

@@ -24,6 +24,9 @@ const (
 	OptionServerJSON      = "server-json"
 	OptionSource          = "source"
 	OptionTimeout         = "timeout"
+	OptionChunks          = "chunks"
+	OptionUploadSize      = "upload-size"
+	OptionDuration        = "duration"
 	OptionSecure          = "secure"
 	OptionNoPreAllocate   = "no-pre-allocate"
 	OptionVersion         = "version"

defs/server.go

@@ -185,7 +185,7 @@ func (s *Server) PingAndJitter(count int) (float64, float64, error) {
 }
 
 // Download performs the actual download test
-func (s *Server) Download(silent bool, useBytes, useMebi bool, requests int) (float64, int, error) {
+func (s *Server) Download(silent bool, useBytes, useMebi bool, requests int, chunks int, duration time.Duration) (float64, int, error) {
 	t := time.Now()
 	defer func() {
 		s.TLog.Logf("Download took %s", time.Now().Sub(t).String())
@@ -210,7 +210,7 @@ func (s *Server) Download(silent bool, useBytes, useMebi bool, requests int) (fl
 		return 0, 0, err
 	}
 	q := req.URL.Query()
-	q.Set("ckSize", strconv.Itoa(downloadChunks))
+	q.Set("ckSize", strconv.Itoa(chunks))
 	req.URL.RawQuery = q.Encode()
 	req.Header.Set("User-Agent", UserAgent)
 	req.Header.Set("Accept-Encoding", "identity")
@@ -261,7 +261,7 @@ func (s *Server) Download(silent bool, useBytes, useMebi bool, requests int) (fl
 		go doDownload()
 		time.Sleep(200 * time.Millisecond)
 	}
-	timeout := time.After(15 * time.Second)
+	timeout := time.After(duration)
 Loop:
 	for {
 		select {
@@ -277,14 +277,15 @@ Loop:
 }
 
 // Upload performs the actual upload test
-func (s *Server) Upload(noPrealloc, silent, useBytes, useMebi bool, requests int) (float64, int, error) {
+func (s *Server) Upload(noPrealloc, silent, useBytes, useMebi bool, requests int, uploadSize int, duration time.Duration) (float64, int, error) {
 	t := time.Now()
 	defer func() {
 		s.TLog.Logf("Upload took %s", time.Now().Sub(t).String())
 	}()
 
 	counter := NewCounter()
 	counter.SetMebi(useMebi)
+	counter.SetUploadSize(uploadSize)
 
 	if noPrealloc {
 		log.Info("Pre-allocation is disabled, performance might be lower!")
@@ -353,7 +354,7 @@ func (s *Server) Upload(noPrealloc, silent, useBytes, useMebi bool, requests int
 		go doUpload()
 		time.Sleep(200 * time.Millisecond)
 	}
-	timeout := time.After(15 * time.Second)
+	timeout := time.After(duration)
 Loop:
 	for {
 		select {

main.go

@@ -135,9 +135,24 @@ func main() {
 			},
 			&cli.IntFlag{
 				Name:  defs.OptionTimeout,
-				Usage: "HTTP `TIMEOUT` in seconds.",
+				Usage: "HTTP `TIMEOUT` in seconds",
 				Value: 15,
 			},
+			&cli.IntFlag{
+				Name:  defs.OptionDuration,
+				Usage: "Upload and download test duration in seconds",
+				Value: 15,
+			},
+			&cli.IntFlag{
+				Name:  defs.OptionChunks,
+				Usage: "Chunks to download from server, chunk size depends on server configuration",
+				Value: 100,
+			},
+			&cli.IntFlag{
+				Name:  defs.OptionUploadSize,
+				Usage: "Size of payload being uploaded in KiB",
+				Value: 1024,
+			},
 			&cli.BoolFlag{
 				Name: defs.OptionSecure,
 				Usage: "Use HTTPS instead of HTTP when communicating with\n" +

speedtest/helper.go

@@ -82,7 +82,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), c.Int(defs.OptionConcurrent))
+				download, br, err := currentServer.Download(silent, c.Bool(defs.OptionBytes), c.Bool(defs.OptionMebiBytes), c.Int(defs.OptionConcurrent), c.Int(defs.OptionChunks), time.Duration(c.Int(defs.OptionDuration))*time.Second)
 				if err != nil {
 					log.Errorf("Failed to get download speed: %s", err)
 					return err
@@ -97,7 +97,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), c.Int(defs.OptionConcurrent))
+				upload, bw, err := currentServer.Upload(c.Bool(defs.OptionNoPreAllocate), silent, c.Bool(defs.OptionBytes), c.Bool(defs.OptionMebiBytes), c.Int(defs.OptionConcurrent), c.Int(defs.OptionUploadSize), time.Duration(c.Int(defs.OptionDuration))*time.Second)
 				if err != nil {
 					log.Errorf("Failed to get upload speed: %s", err)
 					return err