Repository: snowflakedb/gosnowflake
Pull request #982 (2023-12-01 11:38:56)

SNOW-981533 Separate retry strategy for auth endpoints and the remaining ones
### Description
Fixed retry algorithm. Before, only whitelisted endpoints were retried. Now it retries all endpoints. Also, there is a separate backoff logic for authentication endpoints.

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [x] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/982?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`5c79db8`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/5c79db8fc48e0096a65eaaa8b1f8f473ca7205d5?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) 82.78% compared to head [(`1811d73`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/982?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) 80.85%.
> Report is 2 commits behind head on master.

| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/982?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Patch % | Lines |
|---|---|---|
| [retry.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/982?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-cmV0cnkuZ28=) | 96.55% | [1 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/982?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master     #982      +/-   ##
==========================================
- Coverage   82.78%   80.85%   -1.94%     
==========================================
  Files          50       49       -1     
  Lines        9633     9659      +26     
==========================================
- Hits         7975     7810     -165     
- Misses       1658     1849     +191     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/982?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).


Patch:
diff --git a/monitoring.go b/monitoring.go
index 07b17c0..444a7c4 100644
--- a/monitoring.go
+++ b/monitoring.go
@@ -136,7 +136,7 @@ func (sc *snowflakeConn) checkQueryStatus(
 	if tok, _, _ := sc.rest.TokenAccessor.GetTokens(); tok != "" {
 		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, tok)
 	}
-	resultPath := fmt.Sprintf("/monitoring/queries/%s", qid)
+	resultPath := fmt.Sprintf("%s/%s", monitoringQueriesPath, qid)
 	url := sc.rest.getFullURL(resultPath, &param)
 
 	res, err := sc.rest.FuncGet(ctx, sc.rest, url, headers, sc.rest.RequestTimeout)
diff --git a/restful.go b/restful.go
index 777d94d..c92d9c7 100644
--- a/restful.go
+++ b/restful.go
@@ -37,6 +37,7 @@ const (
 	tokenRequestPath         = "/session/token-request"
 	abortRequestPath         = "/queries/v1/abort-request"
 	authenticatorRequestPath = "/session/authenticator-request"
+	monitoringQueriesPath    = "/monitoring/queries"
 	sessionRequestPath       = "/session"
 	heartBeatPath            = "/session/heartbeat"
 )
diff --git a/retry.go b/retry.go
index e46148f..b70b97e 100644
--- a/retry.go
+++ b/retry.go
@@ -20,18 +20,17 @@ import (
 type waitAlgo struct {
 	mutex  *sync.Mutex // required for *rand.Rand usage
 	random *rand.Rand
+	base   time.Duration // base wait time
+	cap    time.Duration // maximum wait time
 }
 
 var random *rand.Rand
 var defaultWaitAlgo *waitAlgo
 
-var endpointsEligibleForRetry = []string{
+var authEndpoints = []string{
 	loginRequestPath,
 	tokenRequestPath,
 	authenticatorRequestPath,
-	queryRequestPath,
-	abortRequestPath,
-	sessionRequestPath,
 }
 
 var clientErrorsStatusCodesEligibleForRetry = []int{
@@ -43,7 +42,7 @@ var clientErrorsStatusCodesEligibleForRetry = []int{
 
 func init() {
 	random = rand.New(rand.NewSource(time.Now().UnixNano()))
-	defaultWaitAlgo = &waitAlgo{mutex: &sync.Mutex{}, random: random}
+	defaultWaitAlgo = &waitAlgo{mutex: &sync.Mutex{}, random: random, base: 5 * time.Second, cap: 160 * time.Second}
 }
 
 const (
@@ -205,12 +204,30 @@ func isQueryRequest(url *url.URL) bool {
 }
 
 // jitter backoff in seconds
-func (w *waitAlgo) calculateWaitBeforeRetry(attempt int, currWaitTime float64) float64 {
+func (w *waitAlgo) calculateWaitBeforeRetryForAuthRequest(attempt int, currWaitTimeDuration time.Duration) time.Duration {
 	w.mutex.Lock()
 	defer w.mutex.Unlock()
-	jitterAmount := w.getJitter(currWaitTime)
-	jitteredSleepTime := chooseRandomFromRange(currWaitTime+jitterAmount, math.Pow(2, float64(attempt))+jitterAmount)
-	return jitteredSleepTime
+	currWaitTimeInSeconds := currWaitTimeDuration.Seconds()
+	jitterAmount := w.getJitter(currWaitTimeInSeconds)
+	jitteredSleepTime := chooseRandomFromRange(currWaitTimeInSeconds+jitterAmount, math.Pow(2, float64(attempt))+jitterAmount)
+	return time.Duration(jitteredSleepTime * float64(time.Second))
+}
+
+func (w *waitAlgo) calculateWaitBeforeRetry(attempt int, sleep time.Duration) time.Duration {
+	w.mutex.Lock()
+	defer w.mutex.Unlock()
+	t := 3*sleep - w.base
+	switch {
+	case t > 0:
+		return durationMin(w.cap, randSecondDuration(t)+w.base)
+	case t < 0:
+		return durationMin(w.cap, randSecondDuration(-t)+3*sleep)
+	}
+	return w.base
+}
+
+func randSecondDuration(n time.Duration) time.Duration {
+	return time.Duration(random.Int63n(int64(n/time.Second))) * time.Second
 }
 
 func (w *waitAlgo) getJitter(currWaitTime float64) float64 {
@@ -284,7 +301,7 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 	totalTimeout := r.timeout
 	logger.WithContext(r.ctx).Infof("retryHTTP.totalTimeout: %v", totalTimeout)
 	retryCounter := 0
-	sleepTime := 1.0 // seconds
+	sleepTime := time.Duration(time.Second)
 	clientStartTime := strconv.FormatInt(r.currentTimeProvider.currentTime(), 10)
 
 	var requestGUIDReplacer requestGUIDReplacer
@@ -324,12 +341,16 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 		}
 		// uses exponential jitter backoff
 		retryCounter++
-		sleepTime = defaultWaitAlgo.calculateWaitBeforeRetry(retryCounter, sleepTime)
+		if isLoginRequest(req) {
+			sleepTime = defaultWaitAlgo.calculateWaitBeforeRetryForAuthRequest(retryCounter, sleepTime)
+		} else {
+			sleepTime = defaultWaitAlgo.calculateWaitBeforeRetry(retryCounter, sleepTime)
+		}
 
 		if totalTimeout > 0 {
 			logger.WithContext(r.ctx).Infof("to timeout: %v", totalTimeout)
 			// if any timeout is set
-			totalTimeout -= time.Duration(sleepTime * float64(time.Second))
+			totalTimeout -= sleepTime
 			if totalTimeout <= 0 || retryCounter > r.maxRetryCount {
 				if err != nil {
 					return nil, err
@@ -360,7 +381,7 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 		logger.WithContext(r.ctx).Infof("sleeping %v. to timeout: %v. retrying", sleepTime, totalTimeout)
 		logger.WithContext(r.ctx).Infof("retry count: %v, retry reason: %v", retryCounter, retryReason)
 
-		await := time.NewTimer(time.Duration(sleepTime * float64(time.Second)))
+		await := time.NewTimer(sleepTime)
 		select {
 		case <-await.C:
 			// retry the request
@@ -378,10 +399,13 @@ func isRetryableError(req *http.Request, res *http.Response, err error) (bool, e
 	if res == nil || req == nil {
 		return false, err
 	}
-	isRetryableURL := contains(endpointsEligibleForRetry, req.URL.Path)
-	return isRetryableURL && isRetryableStatus(res.StatusCode), err
+	return isRetryableStatus(res.StatusCode), err
 }
 
 func isRetryableStatus(statusCode int) bool {
 	return (statusCode >= 500 && statusCode < 600) || contains(clientErrorsStatusCodesEligibleForRetry, statusCode)
 }
+
+func isLoginRequest(req *http.Request) bool {
+	return contains(authEndpoints, req.URL.Path)
+}


Test patch:
diff --git a/retry_test.go b/retry_test.go
index 90c77de..0b280bf 100644
--- a/retry_test.go
+++ b/retry_test.go
@@ -493,12 +493,6 @@ func TestIsRetryable(t *testing.T) {
 			err:      nil,
 			expected: false,
 		},
-		{
-			req:      &http.Request{URL: &url.URL{Path: heartBeatPath}},
-			res:      &http.Response{StatusCode: http.StatusBadRequest},
-			err:      nil,
-			expected: false,
-		},
 		{
 			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
 			res:      &http.Response{StatusCode: http.StatusNotFound},
@@ -525,10 +519,12 @@ func TestIsRetryable(t *testing.T) {
 	}
 
 	for _, tc := range tcs {
-		result, _ := isRetryableError(tc.req, tc.res, tc.err)
-		if result != tc.expected {
-			t.Fatalf("expected %v, got %v; request: %v, response: %v", tc.expected, result, tc.req, tc.res)
-		}
+		t.Run(fmt.Sprintf("req %v, resp %v", tc.req, tc.res), func(t *testing.T) {
+			result, _ := isRetryableError(tc.req, tc.res, tc.err)
+			if result != tc.expected {
+				t.Fatalf("expected %v, got %v; request: %v, response: %v", tc.expected, result, tc.req, tc.res)
+			}
+		})
 	}
 }
 
@@ -605,8 +601,8 @@ func TestCalculateRetryWait(t *testing.T) {
 
 	for _, tc := range tcs {
 		t.Run(fmt.Sprintf("attmept: %v", tc.attempt), func(t *testing.T) {
-			result := defaultWaitAlgo.calculateWaitBeforeRetry(tc.attempt, tc.currWaitTime)
-			assertBetweenE(t, result, tc.minSleepTime, tc.maxSleepTime)
+			result := defaultWaitAlgo.calculateWaitBeforeRetryForAuthRequest(tc.attempt, time.Duration(tc.currWaitTime*float64(time.Second)))
+			assertBetweenE(t, result.Seconds(), tc.minSleepTime, tc.maxSleepTime)
 		})
 	}
 }