Skip to content

Commit e5a11e1

Browse files
committed
Add support for setting per-attempt-timeout
- At the moment, only "complete" requests can be retried (that result in an HTTP status).
1 parent 401eada commit e5a11e1

File tree

5 files changed

+128
-0
lines changed

5 files changed

+128
-0
lines changed

authentication/authentication_option.go

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,26 @@ func WithRetries(maxRetries int, statuses []int) Option {
8383
}
8484
}
8585

86+
// RetryStrategy defines the retry rules that should be followed by the SDK when making requests.
87+
type RetryStrategy struct {
88+
MaxRetries int
89+
Statuses []int
90+
91+
// PerAttemptTimeout can optionally be set to timeout individual API requests.
92+
PerAttemptTimeout time.Duration
93+
}
94+
95+
// WithRetryStrategy configures the management client to only retry under the conditions provided.
96+
func WithRetryStrategy(retryStrategy RetryStrategy) Option {
97+
return func(a *Authentication) {
98+
a.retryStrategy = client.RetryOptions{
99+
MaxRetries: retryStrategy.MaxRetries,
100+
Statuses: retryStrategy.Statuses,
101+
PerAttemptTimeout: retryStrategy.PerAttemptTimeout,
102+
}
103+
}
104+
}
105+
86106
// WithNoRetries configures the management client to only retry under the conditions provided.
87107
func WithNoRetries() Option {
88108
return func(a *Authentication) {

authentication/authentication_test.go

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,6 +450,37 @@ func TestRetries(t *testing.T) {
450450
assert.ErrorIs(t, err, context.Canceled)
451451
assert.Equal(t, 1, i) // 1 request should have been made before the context times out
452452
})
453+
454+
t.Run("Retry strategy", func(t *testing.T) {
455+
i := 0
456+
ctx, cancel := context.WithCancel(context.Background())
457+
458+
h := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
459+
i++
460+
cancel()
461+
w.WriteHeader(http.StatusBadGateway)
462+
})
463+
464+
s := httptest.NewTLSServer(h)
465+
defer s.Close()
466+
467+
a, err := New(
468+
context.Background(),
469+
s.URL,
470+
WithIDTokenSigningAlg("HS256"),
471+
WithClient(s.Client()),
472+
WithRetryStrategy(RetryStrategy{
473+
MaxRetries: 3,
474+
Statuses: []int{http.StatusBadGateway},
475+
PerAttemptTimeout: time.Second,
476+
}),
477+
)
478+
assert.NoError(t, err)
479+
480+
_, err = a.UserInfo(ctx, "123")
481+
assert.ErrorIs(t, err, context.Canceled)
482+
assert.Equal(t, 1, i) // 1 request should have been made before the context times out
483+
})
453484
}
454485

455486
func TestWithClockTolerance(t *testing.T) {

internal/client/client.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@ var DefaultAuth0ClientInfo = &Auth0ClientInfo{
5757
type RetryOptions struct {
5858
MaxRetries int
5959
Statuses []int
60+
61+
// PerAttemptTimeout can optionally be set to timeout individual API requests.
62+
PerAttemptTimeout time.Duration
6063
}
6164

6265
// 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
111114
),
112115
backoffDelay(),
113116
)
117+
tr.PerAttemptTimeout = r.PerAttemptTimeout
114118

115119
return tr
116120
}

management/management_option.go

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package management
33
import (
44
"context"
55
"net/http"
6+
"time"
67

78
"github.com/auth0/go-auth0/internal/client"
89
)
@@ -123,6 +124,26 @@ func WithRetries(maxRetries int, statuses []int) Option {
123124
}
124125
}
125126

127+
// RetryStrategy defines the retry rules that should be followed by the SDK when making requests.
128+
type RetryStrategy struct {
129+
MaxRetries int
130+
Statuses []int
131+
132+
// PerAttemptTimeout can optionally be set to timeout individual API requests.
133+
PerAttemptTimeout time.Duration
134+
}
135+
136+
// WithRetryStrategy configures the management client to only retry under the conditions provided.
137+
func WithRetryStrategy(retryStrategy RetryStrategy) Option {
138+
return func(m *Management) {
139+
m.retryStrategy = client.RetryOptions{
140+
MaxRetries: retryStrategy.MaxRetries,
141+
Statuses: retryStrategy.Statuses,
142+
PerAttemptTimeout: retryStrategy.PerAttemptTimeout,
143+
}
144+
}
145+
}
146+
126147
// WithNoRetries configures the management client to only retry under the conditions provided.
127148
func WithNoRetries() Option {
128149
return func(m *Management) {

management/management_test.go

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"net/http/httptest"
1111
"os"
1212
"runtime"
13+
"sync/atomic"
1314
"testing"
1415
"time"
1516

@@ -529,6 +530,57 @@ func TestRetries(t *testing.T) {
529530
assert.ErrorIs(t, err, context.Canceled)
530531
assert.Equal(t, 1, i) // 1 request should have been made before the context times out
531532
})
533+
534+
t.Run("Retry per request timeout", func(t *testing.T) {
535+
var i atomic.Int64
536+
ctx, cancel := context.WithCancel(context.Background())
537+
defer cancel()
538+
539+
h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
540+
ctx := r.Context()
541+
542+
c := i.Add(1)
543+
t.Log(c)
544+
if c == 2 {
545+
cancel()
546+
}
547+
548+
timer := time.NewTimer(10 * time.Second)
549+
select {
550+
case <-timer.C:
551+
t.Log("completed")
552+
w.WriteHeader(http.StatusOK)
553+
return
554+
case <-ctx.Done():
555+
t.Log("cancelled")
556+
w.WriteHeader(499)
557+
return
558+
}
559+
})
560+
561+
s := httptest.NewServer(h)
562+
defer s.Close()
563+
564+
m, err := New(
565+
s.URL,
566+
WithInsecure(),
567+
WithRetryStrategy(RetryStrategy{
568+
MaxRetries: 10,
569+
Statuses: []int{
570+
http.StatusInternalServerError,
571+
http.StatusBadGateway,
572+
http.StatusServiceUnavailable,
573+
http.StatusGatewayTimeout,
574+
},
575+
PerAttemptTimeout: 5 * time.Millisecond,
576+
}),
577+
)
578+
assert.NoError(t, err)
579+
580+
_, err = m.User.Read(ctx, "123")
581+
assert.ErrorIs(t, err, context.Canceled)
582+
assert.Equal(t, int64(2), i.Load()) // 1 request should have been made before the context times out
583+
})
532584
}
533585

534586
func TestApiCallContextTimeout(t *testing.T) {

0 commit comments

Comments
 (0)