Repository: snowflakedb/gosnowflake
Pull request #919 (2023-10-02 09:03:55)

SNOW-878073 Refactor retry policy to support HTTP 503 & 429
### Description
Please explain the changes you made here.

### Checklist
- [ ] Code compiles correctly
- [ ] Run ``make fmt`` to fix inconsistent formats
- [ ] Run ``make lint`` to get lint errors and fix all of them
- [ ] Created tests which fail without the change (if possible)
- [ ] All tests passing
- [ ] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
> Merging [#919](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) (1835cf9) into [master](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/99921d481830252887ec89396a6ebcddb5856ec4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) (99921d4) will **decrease** coverage by `1.62%`.
> The diff coverage is `98.18%`.

```diff
@@            Coverage Diff             @@
##           master     #919      +/-   ##
==========================================
- Coverage   83.24%   81.62%   -1.62%     
==========================================
  Files          49       48       -1     
  Lines        9340     9323      -17     
==========================================
- Hits         7775     7610     -165     
- Misses       1565     1713     +148     
```


| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Coverage Δ | |
|---|---|---|
| [auth.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXV0aC5nbw==) | `80.38% <100.00%> (ø)` | |
| [authokta.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXV0aG9rdGEuZ28=) | `78.42% <100.00%> (ø)` | |
| [client.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-Y2xpZW50Lmdv) | `100.00% <100.00%> (ø)` | |
| [dsn.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-ZHNuLmdv) | `87.57% <ø> (ø)` | |
| [heartbeat.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-aGVhcnRiZWF0Lmdv) | `69.84% <100.00%> (ø)` | |
| [restful.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-cmVzdGZ1bC5nbw==) | `80.92% <100.00%> (-0.12%)` | :arrow_down: |
| [telemetry.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-dGVsZW1ldHJ5Lmdv) | `97.26% <100.00%> (ø)` | |
| [util.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-dXRpbC5nbw==) | `98.26% <100.00%> (+0.09%)` | :arrow_up: |
| [retry.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-cmV0cnkuZ28=) | `93.65% <97.43%> (+4.81%)` | :arrow_up: |

... and [5 files with indirect coverage changes](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/919/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb)


Patch:
diff --git a/auth.go b/auth.go
index e2f4269..c3894e4 100644
--- a/auth.go
+++ b/auth.go
@@ -226,7 +226,7 @@ func postAuth(
 
 	fullURL := sr.getFullURL(loginRequestPath, params)
 	logger.Infof("full URL: %v", fullURL)
-	resp, err := sr.FuncAuthPost(ctx, client, fullURL, headers, bodyCreator, timeout, true)
+	resp, err := sr.FuncAuthPost(ctx, client, fullURL, headers, bodyCreator, timeout)
 	if err != nil {
 		return nil, err
 	}
@@ -279,6 +279,8 @@ func getHeaders() map[string]string {
 	headers := make(map[string]string)
 	headers[httpHeaderContentType] = headerContentTypeApplicationJSON
 	headers[httpHeaderAccept] = headerAcceptTypeApplicationSnowflake
+	headers[httpClientAppID] = clientType
+	headers[httpClientAppVersion] = SnowflakeGoDriverVersion
 	headers[httpHeaderUserAgent] = userAgent
 	return headers
 }
diff --git a/authokta.go b/authokta.go
index 3e3f3c5..d1a08a4 100644
--- a/authokta.go
+++ b/authokta.go
@@ -216,7 +216,7 @@ func postAuthSAML(
 	fullURL := sr.getFullURL(authenticatorRequestPath, params)
 
 	logger.Infof("fullURL: %v", fullURL)
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true, defaultTimeProvider, nil)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, defaultTimeProvider, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -274,7 +274,7 @@ func postAuthOKTA(
 	if err != nil {
 		return nil, err
 	}
-	resp, err := sr.FuncPost(ctx, sr, targetURL, headers, body, timeout, false, defaultTimeProvider, nil)
+	resp, err := sr.FuncPost(ctx, sr, targetURL, headers, body, timeout, defaultTimeProvider, nil)
 	if err != nil {
 		return nil, err
 	}
diff --git a/client.go b/client.go
index dff709b..14f0e5f 100644
--- a/client.go
+++ b/client.go
@@ -12,7 +12,7 @@ import (
 // InternalClient is implemented by HTTPClient
 type InternalClient interface {
 	Get(context.Context, *url.URL, map[string]string, time.Duration) (*http.Response, error)
-	Post(context.Context, *url.URL, map[string]string, []byte, time.Duration, bool, currentTimeProvider) (*http.Response, error)
+	Post(context.Context, *url.URL, map[string]string, []byte, time.Duration, currentTimeProvider) (*http.Response, error)
 }
 
 type httpClient struct {
@@ -33,7 +33,6 @@ func (cli *httpClient) Post(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	raise4xx bool,
 	currentTimeProvider currentTimeProvider) (*http.Response, error) {
-	return cli.sr.FuncPost(ctx, cli.sr, url, headers, body, timeout, raise4xx, currentTimeProvider, nil)
+	return cli.sr.FuncPost(ctx, cli.sr, url, headers, body, timeout, currentTimeProvider, nil)
 }
diff --git a/connection.go b/connection.go
index a3c04a0..6fcc2c4 100644
--- a/connection.go
+++ b/connection.go
@@ -34,6 +34,8 @@ const (
 	httpHeaderHost             = "Host"
 	httpHeaderValueOctetStream = "application/octet-stream"
 	httpHeaderContentEncoding  = "Content-Encoding"
+	httpClientAppID            = "CLIENT_APP_ID"
+	httpClientAppVersion       = "CLIENT_APP_VERSION"
 )
 
 const (
diff --git a/dsn.go b/dsn.go
index 92fd472..c365e15 100644
--- a/dsn.go
+++ b/dsn.go
@@ -19,7 +19,7 @@ import (
 )
 
 const (
-	defaultClientTimeout          = 900 * time.Second // Timeout for network round trip + read out http response
+	defaultClientTimeout          = 300 * time.Second // Timeout for network round trip + read out http response
 	defaultJWTClientTimeout       = 10 * time.Second  // Timeout for network round trip + read out http response but used for JWT auth
 	defaultLoginTimeout           = 60 * time.Second  // Timeout for retry for login EXCLUDING clientTimeout
 	defaultRequestTimeout         = 0 * time.Second   // Timeout for retry for request EXCLUDING clientTimeout
diff --git a/heartbeat.go b/heartbeat.go
index 9d44bb2..4b40ef6 100644
--- a/heartbeat.go
+++ b/heartbeat.go
@@ -62,7 +62,7 @@ func (hc *heartbeat) heartbeatMain() error {
 
 	fullURL := hc.restful.getFullURL(heartBeatPath, params)
 	timeout := hc.restful.RequestTimeout
-	resp, err := hc.restful.FuncPost(context.Background(), hc.restful, fullURL, headers, nil, timeout, false, defaultTimeProvider, nil)
+	resp, err := hc.restful.FuncPost(context.Background(), hc.restful, fullURL, headers, nil, timeout, defaultTimeProvider, nil)
 	if err != nil {
 		return err
 	}
diff --git a/restful.go b/restful.go
index 6b10dd4..c94e286 100644
--- a/restful.go
+++ b/restful.go
@@ -43,8 +43,8 @@ const (
 
 type (
 	funcGetType      func(context.Context, *snowflakeRestful, *url.URL, map[string]string, time.Duration) (*http.Response, error)
-	funcPostType     func(context.Context, *snowflakeRestful, *url.URL, map[string]string, []byte, time.Duration, bool, currentTimeProvider, *Config) (*http.Response, error)
-	funcAuthPostType func(context.Context, *http.Client, *url.URL, map[string]string, bodyCreatorType, time.Duration, bool) (*http.Response, error)
+	funcPostType     func(context.Context, *snowflakeRestful, *url.URL, map[string]string, []byte, time.Duration, currentTimeProvider, *Config) (*http.Response, error)
+	funcAuthPostType func(context.Context, *http.Client, *url.URL, map[string]string, bodyCreatorType, time.Duration) (*http.Response, error)
 	bodyCreatorType  func() ([]byte, error)
 )
 
@@ -162,14 +162,12 @@ func postRestful(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	raise4XX bool,
 	currentTimeProvider currentTimeProvider,
 	cfg *Config) (
 	*http.Response, error) {
 	return newRetryHTTP(ctx, sr.Client, http.NewRequest, fullURL, headers, timeout, currentTimeProvider, cfg).
 		doPost().
 		setBody(body).
-		doRaise4XX(raise4XX).
 		execute()
 }
 
@@ -189,13 +187,11 @@ func postAuthRestful(
 	fullURL *url.URL,
 	headers map[string]string,
 	bodyCreator bodyCreatorType,
-	timeout time.Duration,
-	raise4XX bool) (
+	timeout time.Duration) (
 	*http.Response, error) {
 	return newRetryHTTP(ctx, client, http.NewRequest, fullURL, headers, timeout, defaultTimeProvider, nil).
 		doPost().
 		setBodyCreator(bodyCreator).
-		doRaise4XX(raise4XX).
 		execute()
 }
 
@@ -243,7 +239,7 @@ func postRestfulQueryHelper(
 
 	var resp *http.Response
 	fullURL := sr.getFullURL(queryRequestPath, params)
-	resp, err = sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true, defaultTimeProvider, cfg)
+	resp, err = sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, defaultTimeProvider, cfg)
 	if err != nil {
 		return nil, err
 	}
@@ -335,7 +331,7 @@ func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Durati
 	token, _, _ := sr.TokenAccessor.GetTokens()
 	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
 
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, nil, 5*time.Second, false, defaultTimeProvider, nil)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, nil, 5*time.Second, defaultTimeProvider, nil)
 	if err != nil {
 		return err
 	}
@@ -394,7 +390,7 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time
 		return err
 	}
 
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqBody, timeout, false, defaultTimeProvider, nil)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqBody, timeout, defaultTimeProvider, nil)
 	if err != nil {
 		return err
 	}
@@ -466,7 +462,7 @@ func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID UUID, time
 		return err
 	}
 
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqByte, timeout, false, defaultTimeProvider, nil)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqByte, timeout, defaultTimeProvider, nil)
 	if err != nil {
 		return err
 	}
diff --git a/retry.go b/retry.go
index 8fa90c8..725a8fc 100644
--- a/retry.go
+++ b/retry.go
@@ -5,23 +5,50 @@ package gosnowflake
 import (
 	"bytes"
 	"context"
-	"crypto/x509"
 	"fmt"
 	"io"
+	"math"
 	"math/rand"
 	"net/http"
 	"net/url"
-	"runtime"
 	"strconv"
 	"strings"
 	"sync"
 	"time"
 )
 
+const (
+	// defaultMaxRetryCount specifies maximum number of subsequent retries
+	defaultMaxRetryCount = 7
+)
+
+type waitAlgo struct {
+	mutex  *sync.Mutex // required for *rand.Rand usage
+	random *rand.Rand
+}
+
 var random *rand.Rand
+var defaultWaitAlgo *waitAlgo
+
+var endpointsEligibleForRetry = []string{
+	loginRequestPath,
+	tokenRequestPath,
+	authenticatorRequestPath,
+	queryRequestPath,
+	abortRequestPath,
+	sessionRequestPath,
+}
+
+var clientErrorsStatusCodesEligibleForRetry = []int{
+	http.StatusTooManyRequests,
+	http.StatusBadRequest,
+	http.StatusMethodNotAllowed,
+	http.StatusRequestTimeout,
+}
 
 func init() {
 	random = rand.New(rand.NewSource(time.Now().UnixNano()))
+	defaultWaitAlgo = &waitAlgo{mutex: &sync.Mutex{}, random: random}
 }
 
 const (
@@ -182,34 +209,28 @@ func isQueryRequest(url *url.URL) bool {
 	return strings.HasPrefix(url.Path, queryRequestPath)
 }
 
-type waitAlgo struct {
-	mutex *sync.Mutex   // required for random.Int63n
-	base  time.Duration // base wait time
-	cap   time.Duration // maximum wait time
-}
-
-func randSecondDuration(n time.Duration) time.Duration {
-	return time.Duration(random.Int63n(int64(n/time.Second))) * time.Second
-}
-
-// decorrelated jitter backoff
-func (w *waitAlgo) decorr(attempt int, sleep time.Duration) time.Duration {
+// jitter backoff in seconds
+func (w *waitAlgo) calculateWaitBeforeRetry(attempt int, currWaitTime float64) float64 {
 	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)
+	var jitterPercentage = 0.5
+	if attempt < 2 {
+		jitterPercentage = 0.25 // to ensure there will be sleep time increase between attempts
 	}
-	return w.base
+	jitterAmount := w.getJitter(currWaitTime, jitterPercentage)
+	jitteredSleepTime := math.Pow(2, float64(attempt)) + jitterAmount
+	return jitteredSleepTime
+}
+
+func (w *waitAlgo) getJitter(currWaitTime float64, jitterPercentage float64) float64 {
+	multiplicationFactor := chooseRandomFromValues(w.random, []int{-1, 1}) // random int from [-1, 1]
+	jitterAmount := jitterPercentage * currWaitTime * float64(multiplicationFactor)
+	return jitterAmount
 }
 
-var defaultWaitAlgo = &waitAlgo{
-	mutex: &sync.Mutex{},
-	base:  5 * time.Second,
-	cap:   160 * time.Second,
+func chooseRandomFromValues[T any](random *rand.Rand, arr []T) T {
+	valIdx := random.Intn(len(arr))
+	return arr[valIdx]
 }
 
 type requestFunc func(method, urlStr string, body io.Reader) (*http.Request, error)
@@ -227,7 +248,6 @@ type retryHTTP struct {
 	headers             map[string]string
 	bodyCreator         bodyCreatorType
 	timeout             time.Duration
-	raise4XX            bool
 	currentTimeProvider currentTimeProvider
 	cfg                 *Config
 }
@@ -249,17 +269,11 @@ func newRetryHTTP(ctx context.Context,
 	instance.headers = headers
 	instance.timeout = timeout
 	instance.bodyCreator = emptyBodyCreator
-	instance.raise4XX = false
 	instance.currentTimeProvider = currentTimeProvider
 	instance.cfg = cfg
 	return &instance
 }
 
-func (r *retryHTTP) doRaise4XX(raise4XX bool) *retryHTTP {
-	r.raise4XX = raise4XX
-	return r
-}
-
 func (r *retryHTTP) doPost() *retryHTTP {
 	r.method = "POST"
 	return r
@@ -281,7 +295,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 := time.Duration(0)
+	sleepTime := 1.0 // seconds
 	clientStartTime := strconv.FormatInt(r.currentTimeProvider.currentTime(), 10)
 
 	var requestGUIDReplacer requestGUIDReplacer
@@ -306,45 +320,37 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 			req.Header.Set(k, v)
 		}
 		res, err = r.client.Do(req)
+		// check if it can retry.
+		retryable, err := isRetryableError(req, res, err)
+		if !retryable {
+			return res, err
+		}
 		if err != nil {
-			// check if it can retry.
-			doExit, err := r.isRetryableError(err)
-			if doExit {
-				return res, err
-			}
-			// cannot just return 4xx and 5xx status as the error can be sporadic. run often helps.
 			logger.WithContext(r.ctx).Warningf(
-				"failed http connection. no response is returned. err: %v. retrying...\n", err)
+				"failed http connection. err: %v. retrying...\n", err)
 		} else {
-			if res.StatusCode == http.StatusOK || r.raise4XX && res != nil && res.StatusCode >= 400 && res.StatusCode < 500 && res.StatusCode != 429 {
-				// exit if success
-				// or
-				// abort connection if raise4XX flag is enabled and the range of HTTP status code are 4XX.
-				// This is currently used for Snowflake login. The caller must generate an error object based on HTTP status.
-				break
-			}
 			logger.WithContext(r.ctx).Warningf(
 				"failed http connection. HTTP Status: %v. retrying...\n", res.StatusCode)
 			res.Body.Close()
 		}
-		// uses decorrelated jitter backoff
-		sleepTime = defaultWaitAlgo.decorr(retryCounter, sleepTime)
+		// uses exponential jitter backoff
+		retryCounter++
+		sleepTime = defaultWaitAlgo.calculateWaitBeforeRetry(retryCounter, sleepTime)
 
 		if totalTimeout > 0 {
 			logger.WithContext(r.ctx).Infof("to timeout: %v", totalTimeout)
 			// if any timeout is set
-			totalTimeout -= sleepTime
-			if totalTimeout <= 0 {
+			totalTimeout -= time.Duration(sleepTime * float64(time.Second))
+			if totalTimeout <= 0 || retryCounter >= defaultMaxRetryCount {
 				if err != nil {
 					return nil, err
 				}
 				if res != nil {
-					return nil, fmt.Errorf("timeout after %s. HTTP Status: %v. Hanging?", r.timeout, res.StatusCode)
+					return nil, fmt.Errorf("timeout after %s and %v retries. HTTP Status: %v. Hanging?", r.timeout, retryCounter, res.StatusCode)
 				}
-				return nil, fmt.Errorf("timeout after %s. Hanging?", r.timeout)
+				return nil, fmt.Errorf("timeout after %s and %v retries. Hanging?", r.timeout, retryCounter)
 			}
 		}
-		retryCounter++
 		if requestGUIDReplacer == nil {
 			requestGUIDReplacer = newRequestGUIDReplace(r.fullURL)
 		}
@@ -365,7 +371,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(sleepTime)
+		await := time.NewTimer(time.Duration(sleepTime * float64(time.Second)))
 		select {
 		case <-await.C:
 			// retry the request
@@ -374,36 +380,19 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 			return res, r.ctx.Err()
 		}
 	}
-	return res, err
 }
 
-func (r *retryHTTP) isRetryableError(err error) (bool, error) {
-	urlError, isURLError := err.(*url.Error)
-	if isURLError {
-		// context cancel or timeout
-		if urlError.Err == context.DeadlineExceeded || urlError.Err == context.Canceled {
-			return true, urlError.Err
-		}
-		if driverError, ok := urlError.Err.(*SnowflakeError); ok {
-			// Certificate Revoked
-			if driverError.Number == ErrOCSPStatusRevoked {
-				return true, err
-			}
-		}
-		if _, ok := urlError.Err.(x509.CertificateInvalidError); ok {
-			// Certificate is invalid
-			return true, err
-		}
-		if _, ok := urlError.Err.(x509.UnknownAuthorityError); ok {
-			// Certificate is self-signed
-			return true, err
-		}
-		errString := urlError.Err.Error()
-		if runtime.GOOS == "darwin" && strings.HasPrefix(errString, "x509:") && strings.HasSuffix(errString, "certificate is expired") {
-			// Certificate is expired
-			return true, err
-		}
-
+func isRetryableError(req *http.Request, res *http.Response, err error) (bool, error) {
+	if err != nil && res == nil { // Failed http connection. Most probably client timeout.
+		return true, err
+	}
+	if res == nil || req == nil {
+		return false, err
 	}
-	return false, err
+	isRetryableURL := contains(endpointsEligibleForRetry, req.URL.Path)
+	return isRetryableURL && isRetryableStatus(res.StatusCode), err
+}
+
+func isRetryableStatus(statusCode int) bool {
+	return (statusCode >= 500 && statusCode < 600) || contains(clientErrorsStatusCodesEligibleForRetry, statusCode)
 }
diff --git a/telemetry.go b/telemetry.go
index 911aba2..1c11c9a 100644
--- a/telemetry.go
+++ b/telemetry.go
@@ -97,7 +97,7 @@ func (st *snowflakeTelemetry) sendBatch() error {
 	}
 	resp, err := st.sr.FuncPost(context.Background(), st.sr,
 		st.sr.getFullURL(telemetryPath, nil), headers, body,
-		defaultTelemetryTimeout, true, defaultTimeProvider, nil)
+		defaultTelemetryTimeout, defaultTimeProvider, nil)
 	if err != nil {
 		logger.Info("failed to upload metrics to telemetry. err: %v", err)
 		return err
diff --git a/util.go b/util.go
index 96f2e9b..1e07bb3 100644
--- a/util.go
+++ b/util.go
@@ -258,3 +258,12 @@ type unixTimeProvider struct {
 func (utp *unixTimeProvider) currentTime() int64 {
 	return time.Now().UnixMilli()
 }
+
+func contains[T comparable](s []T, e T) bool {
+	for _, v := range s {
+		if v == e {
+			return true
+		}
+	}
+	return false
+}


Test patch:
diff --git a/client_test.go b/client_test.go
index 6ef66ef..2e8beaa 100644
--- a/client_test.go
+++ b/client_test.go
@@ -48,7 +48,7 @@ func TestInternalClient(t *testing.T) {
 		t.Fatalf("Expected exactly one GET request, got %v", transport.getRequests)
 	}
 
-	resp, err = internalClient.Post(context.Background(), &url.URL{}, make(map[string]string), make([]byte, 0), 0, false, defaultTimeProvider)
+	resp, err = internalClient.Post(context.Background(), &url.URL{}, make(map[string]string), make([]byte, 0), 0, defaultTimeProvider)
 	if err != nil || resp.StatusCode != 200 {
 		t.Fail()
 	}
diff --git a/driver_ocsp_test.go b/driver_ocsp_test.go
index 8ce109a..4b33174 100644
--- a/driver_ocsp_test.go
+++ b/driver_ocsp_test.go
@@ -634,8 +634,8 @@ func TestOCSPFailClosedResponder404(t *testing.T) {
 	if !ok {
 		t.Fatalf("failed to extract error URL Error: %v", err)
 	}
-	if !strings.Contains(urlErr.Err.Error(), "HTTP Status: 404") {
-		t.Fatalf("the root cause is not  timeout: %v", urlErr.Err)
+	if !strings.Contains(urlErr.Err.Error(), "404 Not Found") {
+		t.Fatalf("the root cause is not timeout: %v", urlErr.Err)
 	}
 }
 
diff --git a/dsn_test.go b/dsn_test.go
index 393eda8..e3e30d6 100644
--- a/dsn_test.go
+++ b/dsn_test.go
@@ -1119,20 +1119,20 @@ func TestDSN(t *testing.T) {
 				User:             "u",
 				Password:         "p",
 				Account:          "a.b.c",
-				ClientTimeout:    300 * time.Second,
+				ClientTimeout:    400 * time.Second,
 				JWTClientTimeout: 60 * time.Second,
 			},
-			dsn: "u:p@a.b.c.snowflakecomputing.com:443?clientTimeout=300&jwtClientTimeout=60&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
+			dsn: "u:p@a.b.c.snowflakecomputing.com:443?clientTimeout=400&jwtClientTimeout=60&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
 		},
 		{
 			cfg: &Config{
 				User:             "u",
 				Password:         "p",
 				Account:          "a.b.c",
-				ClientTimeout:    300 * time.Second,
+				ClientTimeout:    400 * time.Second,
 				JWTExpireTimeout: 30 * time.Second,
 			},
-			dsn: "u:p@a.b.c.snowflakecomputing.com:443?clientTimeout=300&jwtTimeout=30&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
+			dsn: "u:p@a.b.c.snowflakecomputing.com:443?clientTimeout=400&jwtTimeout=30&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
 		},
 		{
 			cfg: &Config{
diff --git a/restful_test.go b/restful_test.go
index 483048b..66b3b30 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -15,63 +15,63 @@ import (
 	"time"
 )
 
-func postTestError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider, _ *Config) (*http.Response, error) {
+func postTestError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ currentTimeProvider, _ *Config) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusOK,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, errors.New("failed to run post method")
 }
 
-func postAuthTestError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ bool) (*http.Response, error) {
+func postAuthTestError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusOK,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, errors.New("failed to run post method")
 }
 
-func postTestSuccessButInvalidJSON(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider, _ *Config) (*http.Response, error) {
+func postTestSuccessButInvalidJSON(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ currentTimeProvider, _ *Config) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusOK,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, nil
 }
 
-func postTestAppBadGatewayError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider, _ *Config) (*http.Response, error) {
+func postTestAppBadGatewayError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ currentTimeProvider, _ *Config) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusBadGateway,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, nil
 }
 
-func postAuthTestAppBadGatewayError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ bool) (*http.Response, error) {
+func postAuthTestAppBadGatewayError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusBadGateway,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, nil
 }
 
-func postTestAppForbiddenError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider, _ *Config) (*http.Response, error) {
+func postTestAppForbiddenError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ currentTimeProvider, _ *Config) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusForbidden,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, nil
 }
 
-func postAuthTestAppForbiddenError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ bool) (*http.Response, error) {
+func postAuthTestAppForbiddenError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusForbidden,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, nil
 }
 
-func postAuthTestAppUnexpectedError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ bool) (*http.Response, error) {
+func postAuthTestAppUnexpectedError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusInsufficientStorage,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, nil
 }
 
-func postTestQueryNotExecuting(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider, _ *Config) (*http.Response, error) {
+func postTestQueryNotExecuting(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ currentTimeProvider, _ *Config) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{
 		Data:    *dd,
@@ -90,7 +90,7 @@ func postTestQueryNotExecuting(_ context.Context, _ *snowflakeRestful, _ *url.UR
 	}, nil
 }
 
-func postTestRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider, _ *Config) (*http.Response, error) {
+func postTestRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ currentTimeProvider, _ *Config) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{
 		Data:    *dd,
@@ -110,7 +110,7 @@ func postTestRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[str
 	}, nil
 }
 
-func postAuthTestAfterRenew(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ bool) (*http.Response, error) {
+func postAuthTestAfterRenew(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{
 		Data:    *dd,
@@ -130,7 +130,7 @@ func postAuthTestAfterRenew(_ context.Context, _ *http.Client, _ *url.URL, _ map
 	}, nil
 }
 
-func postTestAfterRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider, _ *Config) (*http.Response, error) {
+func postTestAfterRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ currentTimeProvider, _ *Config) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{
 		Data:    *dd,
@@ -157,7 +157,7 @@ func cancelTestRetry(ctx context.Context, sr *snowflakeRestful, requestID UUID,
 	if err != nil {
 		return err
 	}
-	resp, err := sr.FuncPost(ctx, sr, &u, getHeaders(), reqByte, timeout, false, defaultTimeProvider, nil)
+	resp, err := sr.FuncPost(ctx, sr, &u, getHeaders(), reqByte, timeout, defaultTimeProvider, nil)
 	if err != nil {
 		return err
 	}
@@ -462,7 +462,7 @@ func TestUnitRenewRestfulSession(t *testing.T) {
 	accessor := getSimpleTokenAccessor()
 	oldToken, oldMasterToken, oldSessionID := "oldtoken", "oldmaster", int64(100)
 	newToken, newMasterToken, newSessionID := "newtoken", "newmaster", int64(200)
-	postTestSuccessWithNewTokens := func(_ context.Context, _ *snowflakeRestful, _ *url.URL, headers map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider, _ *Config) (*http.Response, error) {
+	postTestSuccessWithNewTokens := func(_ context.Context, _ *snowflakeRestful, _ *url.URL, headers map[string]string, _ []byte, _ time.Duration, _ currentTimeProvider, _ *Config) (*http.Response, error) {
 		if headers[headerAuthorizationKey] != fmt.Sprintf(headerSnowflakeToken, oldMasterToken) {
 			t.Fatalf("authorization key doesn't match, %v vs %v", headers[headerAuthorizationKey], fmt.Sprintf(headerSnowflakeToken, oldMasterToken))
 		}
diff --git a/retry_test.go b/retry_test.go
index 397da9e..7a089ec 100644
--- a/retry_test.go
+++ b/retry_test.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -341,10 +342,10 @@ func TestRetryQuerySuccessWithTimeout(t *testing.T) {
 }
 
 func TestRetryQueryFail(t *testing.T) {
-	logger.Info("Retry N times and Fail")
+	logger.Info("Retry N times until there is a timeout and Fail")
 	client := &fakeHTTPClient{
-		cnt:     4,
-		success: false,
+		statusCode: http.StatusTooManyRequests,
+		success:    false,
 	}
 	urlPtr, err := url.Parse("https://fakeaccountretryfail.snowflakecomputing.com:443/queries/v1/query-request?" + requestIDKey)
 	if err != nil {
@@ -352,7 +353,7 @@ func TestRetryQueryFail(t *testing.T) {
 	}
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute()
+		emptyRequest, urlPtr, make(map[string]string), 15*time.Second, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute()
 	if err == nil {
 		t.Fatal("should fail to run retry")
 	}
@@ -470,7 +471,7 @@ func TestLoginRetry429(t *testing.T) {
 	}
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider, nil).doRaise4XX(true).doPost().setBody([]byte{0}).execute() // enable doRaise4XXX
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute() // enable doRaise4XXX
 	if err != nil {
 		t.Fatal("failed to run retry")
 	}
@@ -483,3 +484,83 @@ func TestLoginRetry429(t *testing.T) {
 		t.Fatalf("no retry counter should be attached: %v", retryCountKey)
 	}
 }
+
+func TestIsRetryable(t *testing.T) {
+	tcs := []struct {
+		req      *http.Request
+		res      *http.Response
+		err      error
+		expected bool
+	}{
+		{
+			req:      nil,
+			res:      nil,
+			err:      nil,
+			expected: false,
+		},
+		{
+			req:      nil,
+			res:      &http.Response{StatusCode: http.StatusBadRequest},
+			err:      nil,
+			expected: false,
+		},
+		{
+			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
+			res:      nil,
+			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},
+			expected: false,
+		},
+		{
+			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
+			res:      nil,
+			err:      errUnknownError(),
+			expected: true,
+		},
+		{
+			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
+			res:      &http.Response{StatusCode: http.StatusTooManyRequests},
+			err:      nil,
+			expected: true,
+		},
+		{
+			req:      &http.Request{URL: &url.URL{Path: queryRequestPath}},
+			res:      &http.Response{StatusCode: http.StatusServiceUnavailable},
+			err:      nil,
+			expected: true,
+		},
+	}
+
+	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)
+		}
+	}
+}
+
+func TestExponentialJitterBackoff(t *testing.T) {
+	retryTimes := make([]float64, 10)
+	inputTime := 1.0
+	for i := 0; i < 10; i++ {
+		resultTime := defaultWaitAlgo.calculateWaitBeforeRetry(i+1, inputTime)
+		retryTimes[i] = resultTime
+		inputTime = resultTime
+	}
+
+	for i := 0; i < 9; i++ {
+		if retryTimes[i] >= retryTimes[i+1] {
+			log.Fatalf("expected consequent values to be greater than previous ones; array: %v", retryTimes)
+		}
+	}
+}
diff --git a/telemetry_test.go b/telemetry_test.go
index 5cc1dc7..e3a7cab 100644
--- a/telemetry_test.go
+++ b/telemetry_test.go
@@ -107,7 +107,7 @@ func TestEnableTelemetry(t *testing.T) {
 	})
 }
 
-func funcPostTelemetryRespFail(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider, _ *Config) (*http.Response, error) {
+func funcPostTelemetryRespFail(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ currentTimeProvider, _ *Config) (*http.Response, error) {
 	return nil, errors.New("failed to upload metrics to telemetry")
 }
 
diff --git a/util_test.go b/util_test.go
index 2d8576f..0f10092 100644
--- a/util_test.go
+++ b/util_test.go
@@ -364,3 +364,23 @@ func TestGetFromEnvFailOnMissing(t *testing.T) {
 		t.Error("should report error when there is missing env parameter")
 	}
 }
+
+type tcContains[T comparable] struct {
+	arr      []T
+	e        T
+	expected bool
+}
+
+func TestContains(t *testing.T) {
+	performContainsTestcase(tcContains[int]{[]int{1, 2, 3, 5}, 4, false}, t)
+	performContainsTestcase(tcContains[string]{[]string{"a", "b", "C", "F"}, "C", true}, t)
+	performContainsTestcase(tcContains[int]{[]int{1, 2, 3, 5}, 2, true}, t)
+	performContainsTestcase(tcContains[string]{[]string{"a", "b", "C", "F"}, "f", false}, t)
+}
+
+func performContainsTestcase[S comparable](tc tcContains[S], t *testing.T) {
+	result := contains(tc.arr, tc.e)
+	if result != tc.expected {
+		t.Errorf("contains failed; arr: %v, e: %v, should be %v but was %v", tc.arr, tc.e, tc.expected, result)
+	}
+}