Tweak retry behaviour

In an retryable error situation should not retry if the client method is not considered retryable.
Signed-off-by: Raymond Cypher <raymond.cypher@databricks.com>

Author: rcypher-databricks

internal/client/client.go

@@ -45,8 +45,11 @@ const (
 
 type clientMethod int
 
+//go:generate go run golang.org/x/tools/cmd/stringer -type=clientMethod
+
 const (
-	openSession clientMethod = iota
+	unknown clientMethod = iota
+	openSession
 	closeSession
 	fetchResults
 	getResultSetMetadata
@@ -56,7 +59,9 @@ const (
 	cancelOperation
 )
 
-var nonRetryableClientMethods map[clientMethod]any = map[clientMethod]any{executeStatement: struct{}{}}
+var nonRetryableClientMethods map[clientMethod]any = map[clientMethod]any{
+	executeStatement: struct{}{},
+	unknown:          struct{}{}}
 
 // OpenSession is a wrapper around the thrift operation OpenSession
 // If RecordResults is true, the results will be marshalled to JSON format and written to OpenSession<index>.json
@@ -314,6 +319,10 @@ func SprintGuid(bts []byte) string {
 var retryableStatusCodes = map[int]any{http.StatusTooManyRequests: struct{}{}, http.StatusServiceUnavailable: struct{}{}}
 
 func isRetryableServerResponse(resp *http.Response) bool {
+	if resp == nil {
+		return false
+	}
+
 	_, ok := retryableStatusCodes[resp.StatusCode]
 	return ok
 }
@@ -488,29 +497,24 @@ func RetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, err
 		return false, ctx.Err()
 	}
 
-	if err != nil {
-		if v, ok := err.(*url.Error); ok {
-			s := v.Error()
-			for _, re := range errorRes {
-				if re.MatchString(s) {
-					return false, v
-				}
-			}
+	caller, _ := ctx.Value(ClientMethod).(clientMethod)
+	_, nonRetryableClientMethod := nonRetryableClientMethods[caller]
 
-			if _, ok := v.Err.(x509.UnknownAuthorityError); ok {
-				return false, v
-			}
+	if err != nil {
+		if isRetryableError(err) && !nonRetryableClientMethod {
+			return true, nil
 		}
 
-		// The error is likely recoverable so retry.
-		return true, nil
+		return false, err
 	}
 
-	var checkErr error
-	if resp.StatusCode != http.StatusOK {
-		checkErr = fmt.Errorf("unexpected HTTP status %s", resp.Status)
+	// shouldn't retry on no response or success
+	if resp == nil || resp.StatusCode == http.StatusOK {
+		return false, nil
 	}
 
+	checkErr := fmt.Errorf("unexpected HTTP status %s", resp.Status)
+
 	// 429 Too Many Requests or 503 service unavailable is recoverable. Sometimes the server puts
 	// a Retry-After response header to indicate when the server is
 	// available to start processing request from client.
@@ -523,20 +527,37 @@ func RetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, err
 		return true, dbsqlerrint.NewRetryableError(checkErr, retryAfter)
 	}
 
-	if resp.StatusCode == 0 || (resp.StatusCode >= 500 && resp.StatusCode != http.StatusNotImplemented) {
-		callerAny := ctx.Value(ClientMethod)
-		if caller, ok := callerAny.(clientMethod); ok {
-			if _, noRetry := nonRetryableClientMethods[caller]; !noRetry {
-				return true, checkErr
-			}
-		}
+	if !nonRetryableClientMethod && (resp.StatusCode == 0 || (resp.StatusCode >= 500 && resp.StatusCode != http.StatusNotImplemented)) {
+		return true, checkErr
 	}
 
 	// checkErr will be non-nil if the response code was not StatusOK.
 	// Returning it here ensures that the error handler will be called.
 	return false, checkErr
 }
 
+func isRetryableError(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	if v, ok := err.(*url.Error); ok {
+		s := v.Error()
+		for _, re := range errorRes {
+			if re.MatchString(s) {
+				return false
+			}
+		}
+
+		if _, ok := v.Err.(x509.UnknownAuthorityError); ok {
+			return false
+		}
+	}
+
+	// The error is likely recoverable so retry.
+	return true
+}
+
 func backoff(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
 	// honour the Retry-After header
 	if resp != nil && resp.Header != nil {

internal/client/client_test.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/url"
 	"testing"
+	"time"
 
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/require"
@@ -47,36 +48,76 @@ func TestSprintByteId(t *testing.T) {
 
 func TestRetryPolicy(t *testing.T) {
 
-	t.Run("test handling url error", func(t *testing.T) {
-		urlErr := &url.Error{Err: errors.New("EOF")}
+	t.Run("test handling error", func(t *testing.T) {
 		resp := &http.Response{}
 
-		urlErr.Err = x509.UnknownAuthorityError{}
-		b, err := RetryPolicy(context.Background(), resp, urlErr)
-		require.False(t, b)
-		require.Equal(t, urlErr, err)
+		retryableErrors := []error{
+			&url.Error{Err: errors.New("EOF")},
+			&url.Error{Err: errors.New("other error")},
+			errors.New("non-url error")}
 
+		nonRetryableErrors := []error{&url.Error{Err: x509.UnknownAuthorityError{}}}
 		errMsgs := []string{"stopped after 12 redirects", "unsupported protocol scheme", "certificate is not trusted"}
 		for _, msg := range errMsgs {
-			urlErr.Err = errors.New(msg)
-			b, err := RetryPolicy(context.Background(), resp, urlErr)
+			nonRetryableErrors = append(nonRetryableErrors, &url.Error{Err: errors.New(msg)})
+		}
+
+		// should never retry if context is cancelled/timed out, regardless of client method or error
+		checkCtxCancelled := func(callingContext context.Context, err error) {
+			cancelled, cancel := context.WithCancel(callingContext)
+			cancel()
+			timedOut, cancel2 := context.WithDeadline(callingContext, time.Now().Add(-10*time.Second))
+			defer cancel2()
+
+			// should never retry if context is cancelled/timed out
+			b, checkErr := RetryPolicy(cancelled, resp, err)
+			require.False(t, b)
+			require.Equal(t, cancelled.Err(), checkErr)
+			b, checkErr = RetryPolicy(timedOut, resp, err)
 			require.False(t, b)
-			require.Equal(t, urlErr, err)
+			require.Equal(t, timedOut.Err(), checkErr)
 		}
 
-		urlErr.Err = errors.New("other error")
-		b, err = RetryPolicy(context.Background(), resp, urlErr)
-		require.True(t, b)
-		require.Nil(t, err)
+		for _, err := range retryableErrors {
+			// Should retry only if client method is also retryable
+			for cm := range _clientMethod_index {
+				ctx := context.WithValue(context.Background(), ClientMethod, clientMethod(cm))
+
+				checkCtxCancelled(ctx, err)
+
+				_, nonRetryableClientMethod := nonRetryableClientMethods[clientMethod(cm)]
+				b, checkErr := RetryPolicy(ctx, resp, err)
+				if nonRetryableClientMethod {
+					require.False(t, b)
+					require.Equal(t, err, checkErr)
+				} else {
+					require.True(t, b)
+					require.Nil(t, checkErr)
+				}
+			}
+		}
 
-		b, err = RetryPolicy(context.Background(), resp, errors.New("non-url error"))
-		require.True(t, b)
-		require.Nil(t, err)
+		for _, err := range nonRetryableErrors {
+			// Should non retry regardless of client method
+			for cm := range _clientMethod_index {
+				ctx := context.WithValue(context.Background(), ClientMethod, clientMethod(cm))
+
+				checkCtxCancelled(ctx, err)
 
+				b, checkErr := RetryPolicy(ctx, resp, err)
+				require.False(t, b)
+				require.Equal(t, err, checkErr)
+			}
+		}
 	})
 
 	t.Run("test handling status codes", func(t *testing.T) {
-		resp := &http.Response{}
+		var resp *http.Response
+		retry, err := RetryPolicy(context.Background(), resp, nil)
+		require.False(t, retry)
+		require.Nil(t, err)
+
+		resp = &http.Response{}
 		// 429 nd 503 are always retryable
 		retryableCodes := []int{429, 503}
 		// maybe retryable codes: 0 and >= 500, but not 501