minor backoff refactoring

arriven · Mar 28, 2022 · a805c4a · a805c4a
1 parent 8b7da01
commit a805c4a
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 17 deletions.
diff --git a/src/jobs/http.go b/src/jobs/http.go
@@ -158,9 +158,10 @@ func fastHTTPJob(ctx context.Context, logger *zap.Logger, globalConfig *GlobalCo
 
 		if err := sendFastHTTPRequest(client, req, nil); err != nil {
 			logger.Debug("error sending request", zap.Error(err), zap.Any("args", args))
-			backoffController.Handle(ctx, err)
+			utils.Sleep(ctx, backoffController.Increment().GetTimeout())
 		} else {
 			processedTrafficMonitor.Add(uint64(dataSize))
+			backoffController.Reset()
 		}
 	}
 

diff --git a/src/jobs/rawnet.go b/src/jobs/rawnet.go
@@ -75,7 +75,11 @@ func tcpJob(ctx context.Context, logger *zap.Logger, globalConfig *GlobalConfig,
 
 	for jobConfig.Next(ctx) {
 		err = sendTCP(ctx, logger, jobConfig, trafficMonitor, processedTrafficMonitor)
-		backoffController.Handle(ctx, err)
+		if err != nil {
+			utils.Sleep(ctx, backoffController.Increment().GetTimeout())
+		} else {
+			backoffController.Reset()
+		}
 	}
 
 	return nil, nil

diff --git a/src/utils/backoff.go b/src/utils/backoff.go
@@ -42,7 +42,7 @@ func NewBackoffController(c *BackoffConfig) BackoffController {
 	return BackoffController{BackoffConfig: DefaultBackoffConfig()}
 }
 
-func (c BackoffController) getTimeout() time.Duration {
+func (c BackoffController) GetTimeout() time.Duration {
 	result := c.Timeout
 	for i := 0; i < c.count; i++ {
 		result *= time.Duration(c.Multiplier)
@@ -51,25 +51,17 @@ func (c BackoffController) getTimeout() time.Duration {
 	return result
 }
 
-func (c *BackoffController) increment() {
+func (c *BackoffController) Increment() *BackoffController {
 	if c.count < c.Limit {
 		c.count++
 	}
-}
 
-func (c *BackoffController) reset() {
-	c.count = 0
+	// return pointer to itself for usability
+	return c
 }
 
-func (c *BackoffController) Handle(ctx context.Context, err error) {
-	if err == nil {
-		c.reset()
-
-		return
-	}
-
-	c.increment()
-	Sleep(ctx, c.getTimeout())
+func (c *BackoffController) Reset() {
+	c.count = 0
 }
 
 type Counter struct {

diff --git a/src/utils/countrychecker.go b/src/utils/countrychecker.go
@@ -52,7 +52,7 @@ func CheckCountry(countriesToAvoid []string, strictCountryCheck bool) (bool, str
 		log.Printf("Checking IP address, attempt #%d", counter.iter)
 
 		if country, ip, err = fetchLocationInfo(); err != nil {
-			backoffController.Handle(context.Background(), err)
+			Sleep(context.Background(), backoffController.Increment().GetTimeout())
 		}
 	}