exponential backoff implementation (#452)

src/jobs/base.go

@@ -117,23 +117,11 @@ func Get(t string) Job {
 // BasicJobConfig comment for linter
 type BasicJobConfig struct {
 	IntervalMs int `mapstructure:"interval_ms,omitempty"`
-	Count      int `mapstructure:"count,omitempty"`
-
-	iter int
+	utils.Counter
+	*utils.BackoffConfig
 }
 
 // Next comment for linter
 func (c *BasicJobConfig) Next(ctx context.Context) bool {
-	select {
-	case <-ctx.Done():
-		return false
-	case <-time.After(time.Duration(c.IntervalMs) * time.Millisecond):
-		if c.Count <= 0 {
-			return true
-		}
-
-		c.iter++
-
-		return c.iter <= c.Count
-	}
+	return utils.Sleep(ctx, time.Duration(c.IntervalMs)*time.Millisecond) && c.Counter.Next()
 }

src/jobs/http.go

@@ -24,7 +24,6 @@ package jobs
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"log"
 	"time"
@@ -130,6 +129,8 @@ func fastHTTPJob(ctx context.Context, logger *zap.Logger, globalConfig *GlobalCo
 		return nil, err
 	}
 
+	backoffController := utils.NewBackoffController(jobConfig.BackoffConfig)
+
 	client := http.NewClient(ctx, *clientConfig, logger)
 
 	trafficMonitor := metrics.Default.NewWriter(metrics.Traffic, uuid.New().String())
@@ -156,11 +157,8 @@ func fastHTTPJob(ctx context.Context, logger *zap.Logger, globalConfig *GlobalCo
 		trafficMonitor.Add(uint64(dataSize))
 
 		if err := sendFastHTTPRequest(client, req, nil); err != nil {
-			if errors.Is(err, fasthttp.ErrHostClientRedirectToDifferentScheme) {
-				return nil, err
-			}
-
 			logger.Debug("error sending request", zap.Error(err), zap.Any("args", args))
+			backoffController.Handle(ctx, err)
 		} else {
 			processedTrafficMonitor.Add(uint64(dataSize))
 		}

src/jobs/rawnet.go

@@ -57,6 +57,8 @@ func tcpJob(ctx context.Context, logger *zap.Logger, globalConfig *GlobalConfig,
 		return nil, err
 	}
 
+	backoffController := utils.NewBackoffController(jobConfig.BackoffConfig)
+
 	if globalConfig.ProxyURLs != "" {
 		jobConfig.proxyURLs = templates.ParseAndExecute(logger, globalConfig.ProxyURLs, ctx)
 	}
@@ -72,13 +74,14 @@ func tcpJob(ctx context.Context, logger *zap.Logger, globalConfig *GlobalConfig,
 	}
 
 	for jobConfig.Next(ctx) {
-		sendTCP(ctx, logger, jobConfig, trafficMonitor, processedTrafficMonitor)
+		err = sendTCP(ctx, logger, jobConfig, trafficMonitor, processedTrafficMonitor)
+		backoffController.Handle(ctx, err)
 	}
 
 	return nil, nil
 }
 
-func sendTCP(ctx context.Context, logger *zap.Logger, jobConfig *rawnetConfig, trafficMonitor, processedTrafficMonitor *metrics.Writer) {
+func sendTCP(ctx context.Context, logger *zap.Logger, jobConfig *rawnetConfig, trafficMonitor, processedTrafficMonitor *metrics.Writer) error {
 	// track sending of SYN packet
 	trafficMonitor.Add(packetgen.TCPHeaderSize + packetgen.IPHeaderSize)
 
@@ -87,7 +90,7 @@ func sendTCP(ctx context.Context, logger *zap.Logger, jobConfig *rawnetConfig, t
 		logger.Debug("error connecting via tcp", zap.String("addr", jobConfig.addr), zap.Error(err))
 		metrics.IncRawnetTCP(jobConfig.addr, metrics.StatusFail)
 
-		return
+		return err
 	}
 
 	defer conn.Close()
@@ -107,12 +110,14 @@ func sendTCP(ctx context.Context, logger *zap.Logger, jobConfig *rawnetConfig, t
 		if err != nil {
 			metrics.IncRawnetTCP(jobConfig.addr, metrics.StatusFail)
 
-			return
+			return err
 		}
 
 		processedTrafficMonitor.Add(uint64(n))
 		metrics.IncRawnetTCP(jobConfig.addr, metrics.StatusSuccess)
 	}
+
+	return nil
 }
 
 func udpJob(ctx context.Context, logger *zap.Logger, _ *GlobalConfig, args Args) (data interface{}, err error) {

src/utils/backoff.go

@@ -0,0 +1,89 @@
+package utils
+
+import (
+	"context"
+	"time"
+)
+
+func Sleep(ctx context.Context, t time.Duration) bool {
+	select {
+	case <-time.After(t):
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}
+
+type BackoffConfig struct {
+	Multiplier int           `mapstructure:"backoff_multiplier"`
+	Limit      int           `mapstructure:"backoff_limit"`
+	Timeout    time.Duration `mapstructure:"backoff_timeout"`
+}
+
+func DefaultBackoffConfig() BackoffConfig {
+	const (
+		defaultMultiplier = 10
+		defaultLimit      = 6
+	)
+
+	return BackoffConfig{Multiplier: defaultMultiplier, Limit: defaultLimit, Timeout: time.Microsecond}
+}
+
+type BackoffController struct {
+	BackoffConfig
+	count int
+}
+
+func NewBackoffController(c *BackoffConfig) BackoffController {
+	if c != nil {
+		return BackoffController{BackoffConfig: *c}
+	}
+
+	return BackoffController{BackoffConfig: DefaultBackoffConfig()}
+}
+
+func (c BackoffController) getTimeout() time.Duration {
+	result := c.Timeout
+	for i := 0; i < c.count; i++ {
+		result *= time.Duration(c.Multiplier)
+	}
+
+	return result
+}
+
+func (c *BackoffController) increment() {
+	if c.count < c.Limit {
+		c.count++
+	}
+}
+
+func (c *BackoffController) reset() {
+	c.count = 0
+}
+
+func (c *BackoffController) Handle(ctx context.Context, err error) {
+	if err == nil {
+		c.reset()
+
+		return
+	}
+
+	c.increment()
+	Sleep(ctx, c.getTimeout())
+}
+
+type Counter struct {
+	Count int `mapstructure:"count,omitempty"`
+
+	iter int
+}
+
+func (c *Counter) Next() bool {
+	if c.Count <= 0 {
+		return true
+	}
+
+	c.iter++
+
+	return c.iter <= c.Count
+}

src/utils/countrychecker.go

@@ -40,29 +40,30 @@ func CheckCountryOrFail(cfg *CountryCheckerConfig) string {
 func CheckCountry(countriesToAvoid []string, strictCountryCheck bool) (bool, string) {
 	const maxFetchRetries = 3
 
-	var country, ip string
-
-	for retries := 1; ; retries++ {
-		log.Printf("Checking IP address, attempt #%d", retries)
+	var (
+		country, ip string
+		err         error
+	)
 
-		var err error
-		if country, ip, err = fetchLocationInfo(); err != nil {
-			if retries < maxFetchRetries {
-				time.Sleep(time.Second)
+	counter := Counter{Count: maxFetchRetries}
+	backoffController := NewBackoffController(nil)
 
-				continue
-			}
+	for counter.Next() {
+		log.Printf("Checking IP address, attempt #%d", counter.iter)
 
-			if strictCountryCheck {
-				log.Printf("Failed to check the country info in %d attempts while in strict mode", maxFetchRetries)
+		if country, ip, err = fetchLocationInfo(); err != nil {
+			backoffController.Handle(context.Background(), err)
+		}
+	}
 
-				return false, ""
-			}
+	if err != nil {
+		if strictCountryCheck {
+			log.Printf("Failed to check the country info in %d attempts while in strict mode", maxFetchRetries)
 
-			return true, ""
+			return false, ""
 		}
 
-		break
+		return true, ""
 	}
 
 	log.Printf("Current country: %s (%s)", country, ip)