From 8daba8aa19e2c504ad62bc76cc08f0c47d6f30f7 Mon Sep 17 00:00:00 2001
From: Phil Bracikowski <13472206+philjb@users.noreply.github.com>
Date: Wed, 25 Aug 2021 15:43:06 -0700
Subject: [PATCH] chore(influxv2plugin): Increase accepted retry-after header
 values. (#9619)

---
 plugins/outputs/influxdb_v2/http.go           | 15 ++++----
 .../outputs/influxdb_v2/http_internal_test.go | 36 ++++++++++++++++---
 2 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/plugins/outputs/influxdb_v2/http.go b/plugins/outputs/influxdb_v2/http.go
index e62919cf4..e8df4da7d 100644
--- a/plugins/outputs/influxdb_v2/http.go
+++ b/plugins/outputs/influxdb_v2/http.go
@@ -36,8 +36,9 @@ func (e APIError) Error() string {
 }
 
 const (
-	defaultRequestTimeout = time.Second * 5
-	defaultMaxWait        = 60 // seconds
+	defaultRequestTimeout           = time.Second * 5
+	defaultMaxWaitSeconds           = 60
+	defaultMaxWaitRetryAfterSeconds = 10 * 60
 )
 
 type HTTPConfig struct {
@@ -306,8 +307,9 @@ func (c *httpClient) writeBatch(ctx context.Context, bucket string, metrics []te
 // retryDuration takes the longer of the Retry-After header and our own back-off calculation
 func (c *httpClient) getRetryDuration(headers http.Header) time.Duration {
 	// basic exponential backoff (x^2)/40 (denominator to widen the slope)
-	// at 40 denominator, it'll take 35 retries to hit the max defaultMaxWait of 30s
+	// at 40 denominator, it'll take 49 retries to hit the max defaultMaxWait of 60s
 	backoff := math.Pow(float64(c.retryCount), 2) / 40
+	backoff = math.Min(backoff, defaultMaxWaitSeconds)
 
 	// get any value from the header, if available
 	retryAfterHeader := float64(0)
@@ -319,11 +321,12 @@ func (c *httpClient) getRetryDuration(headers http.Header) time.Duration {
 			// there was a value but we couldn't parse it? guess minimum 10 sec
 			retryAfterHeader = 10
 		}
+		// protect against excessively large retry-after
+		retryAfterHeader = math.Min(retryAfterHeader, defaultMaxWaitRetryAfterSeconds)
 	}
-	// take the highest value from both, but not over the max wait.
+	// take the highest value of backoff and retry-after.
 	retry := math.Max(backoff, retryAfterHeader)
-	retry = math.Min(retry, defaultMaxWait)
-	return time.Duration(retry) * time.Second
+	return time.Duration(retry*1000) * time.Millisecond
 }
 
 func (c *httpClient) makeWriteRequest(url string, body io.Reader) (*http.Request, error) {
diff --git a/plugins/outputs/influxdb_v2/http_internal_test.go b/plugins/outputs/influxdb_v2/http_internal_test.go
index 2ff4990fa..10e2a4e13 100644
--- a/plugins/outputs/influxdb_v2/http_internal_test.go
+++ b/plugins/outputs/influxdb_v2/http_internal_test.go
@@ -56,12 +56,12 @@ func TestExponentialBackoffCalculation(t *testing.T) {
 		expected   time.Duration
 	}{
 		{retryCount: 0, expected: 0},
-		{retryCount: 1, expected: 0},
-		{retryCount: 5, expected: 0},
-		{retryCount: 10, expected: 2 * time.Second},
-		{retryCount: 30, expected: 22 * time.Second},
+		{retryCount: 1, expected: 25 * time.Millisecond},
+		{retryCount: 5, expected: 625 * time.Millisecond},
+		{retryCount: 10, expected: 2500 * time.Millisecond},
+		{retryCount: 30, expected: 22500 * time.Millisecond},
 		{retryCount: 40, expected: 40 * time.Second},
-		{retryCount: 50, expected: 60 * time.Second},
+		{retryCount: 50, expected: 60 * time.Second}, // max hit
 		{retryCount: 100, expected: 60 * time.Second},
 		{retryCount: 1000, expected: 60 * time.Second},
 	}
@@ -72,3 +72,29 @@ func TestExponentialBackoffCalculation(t *testing.T) {
 		})
 	}
 }
+
+func TestExponentialBackoffCalculationWithRetryAfter(t *testing.T) {
+	c := &httpClient{}
+	tests := []struct {
+		retryCount int
+		retryAfter string
+		expected   time.Duration
+	}{
+		{retryCount: 0, retryAfter: "0", expected: 0},
+		{retryCount: 0, retryAfter: "10", expected: 10 * time.Second},
+		{retryCount: 0, retryAfter: "60", expected: 60 * time.Second},
+		{retryCount: 0, retryAfter: "600", expected: 600 * time.Second},
+		{retryCount: 0, retryAfter: "601", expected: 600 * time.Second}, // max hit
+		{retryCount: 40, retryAfter: "39", expected: 40 * time.Second},  // retryCount wins
+		{retryCount: 40, retryAfter: "41", expected: 41 * time.Second},  // retryAfter wins
+		{retryCount: 100, retryAfter: "100", expected: 100 * time.Second},
+	}
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("%d_retries", test.retryCount), func(t *testing.T) {
+			c.retryCount = test.retryCount
+			hdr := http.Header{}
+			hdr.Add("Retry-After", test.retryAfter)
+			require.EqualValues(t, test.expected, c.getRetryDuration(hdr))
+		})
+	}
+}