Add support for setting per-attempt-timeout

- At the moment, only "complete" requests can be retried (that result in an HTTP status).

Author: pete-woods

authentication/authentication_option.go

@@ -83,6 +83,26 @@ func WithRetries(maxRetries int, statuses []int) Option {
 	}
 }
 
+// RetryStrategy defines the retry rules that should be followed by the SDK when making requests.
+type RetryStrategy struct {
+	MaxRetries int
+	Statuses   []int
+
+	// PerAttemptTimeout can optionally be set to timeout individual API requests.
+	PerAttemptTimeout time.Duration
+}
+
+// WithRetryStrategy configures the management client to only retry under the conditions provided.
+func WithRetryStrategy(retryStrategy RetryStrategy) Option {
+	return func(a *Authentication) {
+		a.retryStrategy = client.RetryOptions{
+			MaxRetries:        retryStrategy.MaxRetries,
+			Statuses:          retryStrategy.Statuses,
+			PerAttemptTimeout: retryStrategy.PerAttemptTimeout,
+		}
+	}
+}
+
 // WithNoRetries configures the management client to only retry under the conditions provided.
 func WithNoRetries() Option {
 	return func(a *Authentication) {

authentication/authentication_test.go

@@ -11,6 +11,7 @@ import (
 	"net/url"
 	"os"
 	"runtime"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -450,6 +451,90 @@ func TestRetries(t *testing.T) {
 		assert.ErrorIs(t, err, context.Canceled)
 		assert.Equal(t, 1, i) // 1 request should have been made before the context times out
 	})
+
+	t.Run("Retry strategy", func(t *testing.T) {
+		i := 0
+		ctx, cancel := context.WithCancel(context.Background())
+
+		h := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			i++
+			cancel()
+			w.WriteHeader(http.StatusBadGateway)
+		})
+
+		s := httptest.NewTLSServer(h)
+		defer s.Close()
+
+		a, err := New(
+			context.Background(),
+			s.URL,
+			WithIDTokenSigningAlg("HS256"),
+			WithClient(s.Client()),
+			WithRetryStrategy(RetryStrategy{
+				MaxRetries:        3,
+				Statuses:          []int{http.StatusBadGateway},
+				PerAttemptTimeout: time.Second,
+			}),
+		)
+		assert.NoError(t, err)
+
+		_, err = a.UserInfo(ctx, "123")
+		assert.ErrorIs(t, err, context.Canceled)
+		assert.Equal(t, 1, i) // 1 request should have been made before the context times out
+	})
+
+	t.Run("Retry per request timeout", func(t *testing.T) {
+		var i atomic.Int64
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+
+			c := i.Add(1)
+			t.Log(c)
+			if c == 2 {
+				cancel()
+			}
+
+			timer := time.NewTimer(10 * time.Second)
+			select {
+			case <-timer.C:
+				t.Log("completed")
+				w.WriteHeader(http.StatusOK)
+				return
+			case <-ctx.Done():
+				t.Log("cancelled")
+				w.WriteHeader(499)
+				return
+			}
+		})
+
+		s := httptest.NewTLSServer(h)
+		defer s.Close()
+
+		m, err := New(
+			context.Background(),
+			s.URL,
+			WithIDTokenSigningAlg("HS256"),
+			WithClient(s.Client()),
+			WithRetryStrategy(RetryStrategy{
+				MaxRetries: 10,
+				Statuses: []int{
+					http.StatusInternalServerError,
+					http.StatusBadGateway,
+					http.StatusServiceUnavailable,
+					http.StatusGatewayTimeout,
+				},
+				PerAttemptTimeout: 5 * time.Millisecond,
+			}),
+		)
+		assert.NoError(t, err)
+
+		_, err = m.UserInfo(ctx, "123")
+		assert.ErrorIs(t, err, context.Canceled)
+		assert.Equal(t, int64(2), i.Load()) // 1 request should have been made before the context times out
+	})
 }
 
 func TestWithClockTolerance(t *testing.T) {

internal/client/client.go

@@ -57,6 +57,9 @@ var DefaultAuth0ClientInfo = &Auth0ClientInfo{
 type RetryOptions struct {
 	MaxRetries int
 	Statuses   []int
+
+	// PerAttemptTimeout can optionally be set to timeout individual API requests.
+	PerAttemptTimeout time.Duration
 }
 
 // IsEmpty checks whether the provided Auth0ClientInfo data is nil or has no data to allow
@@ -111,6 +114,7 @@ func RetriesTransport(base http.RoundTripper, r RetryOptions) http.RoundTripper
 		),
 		backoffDelay(),
 	)
+	tr.PerAttemptTimeout = r.PerAttemptTimeout
 
 	return tr
 }

management/management.gen.go

@@ -3,6 +3,7 @@ package management
 import (
 	"context"
 	"net/http"
+	"time"
 
 	"github.com/auth0/go-auth0/internal/client"
 )
@@ -123,6 +124,26 @@ func WithRetries(maxRetries int, statuses []int) Option {
 	}
 }
 
+// RetryStrategy defines the retry rules that should be followed by the SDK when making requests.
+type RetryStrategy struct {
+	MaxRetries int
+	Statuses   []int
+
+	// PerAttemptTimeout can optionally be set to timeout individual API requests.
+	PerAttemptTimeout time.Duration
+}
+
+// WithRetryStrategy configures the management client to only retry under the conditions provided.
+func WithRetryStrategy(retryStrategy RetryStrategy) Option {
+	return func(m *Management) {
+		m.retryStrategy = client.RetryOptions{
+			MaxRetries:        retryStrategy.MaxRetries,
+			Statuses:          retryStrategy.Statuses,
+			PerAttemptTimeout: retryStrategy.PerAttemptTimeout,
+		}
+	}
+}
+
 // WithNoRetries configures the management client to only retry under the conditions provided.
 func WithNoRetries() Option {
 	return func(m *Management) {

management/management.gen_test.go

@@ -10,6 +10,7 @@ import (
 	"net/http/httptest"
 	"os"
 	"runtime"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -529,6 +530,57 @@ func TestRetries(t *testing.T) {
 		assert.ErrorIs(t, err, context.Canceled)
 		assert.Equal(t, 1, i) // 1 request should have been made before the context times out
 	})
+
+	t.Run("Retry per request timeout", func(t *testing.T) {
+		var i atomic.Int64
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+
+			c := i.Add(1)
+			t.Log(c)
+			if c == 2 {
+				cancel()
+			}
+
+			timer := time.NewTimer(10 * time.Second)
+			select {
+			case <-timer.C:
+				t.Log("completed")
+				w.WriteHeader(http.StatusOK)
+				return
+			case <-ctx.Done():
+				t.Log("cancelled")
+				w.WriteHeader(499)
+				return
+			}
+		})
+
+		s := httptest.NewServer(h)
+		defer s.Close()
+
+		m, err := New(
+			s.URL,
+			WithInsecure(),
+			WithRetryStrategy(RetryStrategy{
+				MaxRetries: 10,
+				Statuses: []int{
+					http.StatusInternalServerError,
+					http.StatusBadGateway,
+					http.StatusServiceUnavailable,
+					http.StatusGatewayTimeout,
+				},
+				PerAttemptTimeout: 5 * time.Millisecond,
+			}),
+		)
+		assert.NoError(t, err)
+
+		_, err = m.User.Read(ctx, "123")
+		assert.ErrorIs(t, err, context.Canceled)
+		assert.Equal(t, int64(2), i.Load()) // 1 request should have been made before the context times out
+	})
 }
 
 func TestApiCallContextTimeout(t *testing.T) {