From e82768ff8064a95691301e60b5026c6873b2be77 Mon Sep 17 00:00:00 2001
From: Richard Patel <terorie@alphakevin.club>
Date: Tue, 27 Nov 2018 19:47:30 +0100
Subject: [PATCH] Wait time control in config

---
 config.go  | 10 +++++++++-
 config.yml |  7 +++++++
 main.go    |  4 ++--
 server.go  |  5 +++--
 4 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/config.go b/config.go
index fea0f1c..a9f5771 100644
--- a/config.go
+++ b/config.go
@@ -33,7 +33,11 @@ const (
 	ConfToken      = "server.token"
 	ConfServerTimeout = "server.timeout"
 	ConfRecheck    = "server.recheck"
+	ConfCooldown   = "server.cooldown"
 	ConfChunkSize  = "server.upload_chunk"
+	ConfUploadRetries = "server.upload_retries"
+	ConfUploadRetryInterval = "server.upload_retry_interval"
+
 	ConfTasks      = "crawl.tasks"
 	ConfRetries    = "crawl.retries"
 	ConfWorkers    = "crawl.connections"
@@ -41,6 +45,7 @@ const (
 	ConfDialTimeout = "crawl.dial_timeout"
 	ConfTimeout    = "crawl.timeout"
 	ConfJobBufferSize = "crawl.job_buffer"
+
 	ConfCrawlStats = "output.crawl_stats"
 	ConfAllocStats = "output.resource_stats"
 	ConfVerbose    = "output.verbose"
@@ -54,7 +59,7 @@ func prepareConfig() {
 	viper.SetDefault(ConfTasks, 3)
 	viper.SetDefault(ConfUserAgent, "")
 	viper.SetDefault(ConfDialTimeout, 10 * time.Second)
-	viper.SetDefault(ConfTimeout, 30 * time.Second)
+	viper.SetDefault(ConfTimeout, 60 * time.Second)
 	viper.SetDefault(ConfJobBufferSize, 5000)
 	viper.SetDefault(ConfCrawlStats, 3 * time.Second)
 	viper.SetDefault(ConfAllocStats, 0)
@@ -62,7 +67,10 @@ func prepareConfig() {
 	viper.SetDefault(ConfPrintHTTP, false)
 	viper.SetDefault(ConfLogFile, "")
 	viper.SetDefault(ConfRecheck, 3 * time.Second)
+	viper.SetDefault(ConfCooldown, 30 * time.Second)
 	viper.SetDefault(ConfChunkSize, "1 MB")
+	viper.SetDefault(ConfUploadRetries, 10)
+	viper.SetDefault(ConfUploadRetryInterval, 30 * time.Second)
 }
 
 func readConfig() {
diff --git a/config.yml b/config.yml
index 13e7ae3..d1ecaa3 100644
--- a/config.yml
+++ b/config.yml
@@ -15,10 +15,17 @@ server:
   # between /task/get requests to the server.
   recheck: 1s
 
+  # Time to wait after receiving an error
+  # from the server. Doesn't apply to uploads.
+  cooldown: 30s
+
   # Upload chunk size
   # If the value is too high, the upload fails.
   upload_chunk: 1 MB
 
+  upload_retries: 10
+  upload_retry_interval: 30s
+
 # Log output settings
 output:
   # Crawl statistics
diff --git a/main.go b/main.go
index 331d38b..2b407e2 100644
--- a/main.go
+++ b/main.go
@@ -84,7 +84,7 @@ func cmdBase(_ *cli.Context) error {
 			if err != nil {
 				logrus.WithError(err).
 					Error("Failed to get new task")
-				time.Sleep(30 * time.Second)
+				time.Sleep(viper.GetDuration(ConfCooldown))
 				continue
 			}
 			if t == nil {
@@ -111,7 +111,7 @@ func cmdBase(_ *cli.Context) error {
 			} else if err != nil {
 				logrus.WithError(err).
 					Error("Failed to get new task")
-				time.Sleep(30 * time.Second)
+				time.Sleep(viper.GetDuration(ConfCooldown))
 				continue
 			}
 			ScheduleTask(inRemotes, t, &baseUri)
diff --git a/server.go b/server.go
index e4b5f20..abcf3aa 100644
--- a/server.go
+++ b/server.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/sirupsen/logrus"
+	"github.com/spf13/viper"
 	"io"
 	"mime/multipart"
 	"net/http"
@@ -102,10 +103,10 @@ func uploadChunks(websiteId uint64, f *os.File) error {
 
 		multi.Close()
 
-		for retries := 0; retries < 10; retries++ {
+		for retries := 0; retries < viper.GetInt(ConfUploadRetries); retries++ {
 			if retries > 0 {
 				// Error occurred, retry upload
-				time.Sleep(30 * time.Second)
+				time.Sleep(viper.GetDuration(ConfUploadRetryInterval))
 			}
 
 			req, err := http.NewRequest(