From 67726598a766e72a82221b01d119b3787fd544c2 Mon Sep 17 00:00:00 2001
From: n33pm <12273891+n33pm@users.noreply.github.com>
Date: Mon, 10 Jun 2024 13:51:50 +0200
Subject: [PATCH] feat(): update parallel connections and duration settings

---
 download.go |  2 +-
 main.go     | 36 ++++++++++++++++++++++++++++++++----
 upload.go   |  2 +-
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/download.go b/download.go
index b239281..1d95f68 100644
--- a/download.go
+++ b/download.go
@@ -42,7 +42,7 @@ func startDownloadTest(targets []Target) float64 {
 		totalTime += result.TotalTime
 	}
 
-	downloadSpeed := (float64(totalSize * uint64(time.Second))) / float64(totalTime) * parallelConnections
+	downloadSpeed := (float64(totalSize * uint64(time.Second))) / float64(totalTime) * float64(parallelConnections)
 
 	log.Debug().Dur("totalTime", totalTime).Uint64("totalSize", totalSize).Float64("speed", downloadSpeed/1e3/1e3*8).Msg("DownloadSpeed")
 
diff --git a/main.go b/main.go
index 00d31cb..83e799f 100644
--- a/main.go
+++ b/main.go
@@ -7,14 +7,12 @@ import (
 	"net/http"
 	"os"
 	"regexp"
+	"strconv"
 	"time"
 )
 
 const (
-	downloadDuration    = 30 * time.Second
-	uploadDuration      = 30 * time.Second
-	uploadBufferSize    = 26 * 1024 * 1024
-	parallelConnections = 5
+	uploadBufferSize = 26 * 1024 * 1024
 )
 
 var (
@@ -23,6 +21,11 @@ var (
 
 	commit  = ""
 	version = "dev"
+
+	parallelConnections = 5
+
+	downloadDuration = 30 * time.Second
+	uploadDuration   = 30 * time.Second
 )
 
 type APIResponse struct {
@@ -63,15 +66,40 @@ func init() {
 		zerolog.SetGlobalLevel(zerolog.InfoLevel)
 	}
 
+	// read env vars parallelConnections
+	if os.Getenv("PARALLEL_CONNECTIONS") != "" {
+		i, err := strconv.Atoi(os.Getenv("PARALLEL_CONNECTIONS"))
+		if err != nil {
+			log.Fatal().Err(err).Msg("Error reading PARALLEL_CONNECTIONS")
+		}
+		parallelConnections = i
+	}
+	if os.Getenv("DOWNLOAD_DURATION_SEC") != "" {
+		i, err := strconv.Atoi(os.Getenv("DOWNLOAD_DURATION_SEC"))
+		if err != nil {
+			log.Fatal().Err(err).Msg("Error reading DOWNLOAD_DURATION_SEC")
+		}
+		downloadDuration = time.Duration(i) * time.Second
+	}
+	if os.Getenv("UPLOAD_DURATION_SEC") != "" {
+		i, err := strconv.Atoi(os.Getenv("UPLOAD_DURATION_SEC"))
+		if err != nil {
+			log.Fatal().Err(err).Msg("Error reading UPLOAD_DURATION_SEC")
+		}
+		uploadDuration = time.Duration(i) * time.Second
+	}
+
 }
 
 func main() {
 
 	log.Info().Str("version", version).Str("commit", commit).Msg("Starting fast-speedtest")
+	log.Info().Int("parallelConnections", parallelConnections).Float64("downloadDurationSec", downloadDuration.Seconds()).Float64("uploadDurationSec", uploadDuration.Seconds()).Msg("Config")
 
 	//startTest()
 	http.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
 		downloadSpeed, uploadSpeed := startTest()
+		log.Info().Float64("downloadSpeed", downloadSpeed).Float64("uploadSpeed", uploadSpeed).Msg("Speedtest results")
 		w.Header().Add("content-type", "text/plain")
 		w.Write([]byte("# TYPE speedtest_bits_per_second gauge\n"))
 		w.Write([]byte("# HELP speedtest_bits_per_second Speed measured against fast.com\n"))
diff --git a/upload.go b/upload.go
index 1ddacfc..36af753 100644
--- a/upload.go
+++ b/upload.go
@@ -46,7 +46,7 @@ func startUploadTest(targets []Target) float64 {
 		totalTime += result.TotalTime
 	}
 
-	speed := (float64(totalSize * uint64(time.Second))) / float64(totalTime) * parallelConnections
+	speed := (float64(totalSize * uint64(time.Second))) / float64(totalTime) * float64(parallelConnections)
 
 	log.Debug().Dur("totalTime", totalTime).Uint64("totalSize", totalSize).Float64("speed", speed/1e3/1e3*8).Msg("UploadSpeed")