1package retry
2
3import (
4	"time"
5
6	"github.com/aws/aws-sdk-go-v2/aws"
7)
8
9// AddWithErrorCodes returns a Retryer with additional error codes considered
10// for determining if the error should be retried.
11func AddWithErrorCodes(r aws.Retryer, codes ...string) aws.Retryer {
12	retryable := &RetryableErrorCode{
13		Codes: map[string]struct{}{},
14	}
15	for _, c := range codes {
16		retryable.Codes[c] = struct{}{}
17	}
18
19	return &withIsErrorRetryable{
20		Retryer:   r,
21		Retryable: retryable,
22	}
23}
24
25type withIsErrorRetryable struct {
26	aws.Retryer
27	Retryable IsErrorRetryable
28}
29
30func (r *withIsErrorRetryable) IsErrorRetryable(err error) bool {
31	if v := r.Retryable.IsErrorRetryable(err); v != aws.UnknownTernary {
32		return v.Bool()
33	}
34	return r.Retryer.IsErrorRetryable(err)
35}
36
37// AddWithMaxAttempts returns a Retryer with MaxAttempts set to the value
38// specified.
39func AddWithMaxAttempts(r aws.Retryer, max int) aws.Retryer {
40	return &withMaxAttempts{
41		Retryer: r,
42		Max:     max,
43	}
44}
45
46type withMaxAttempts struct {
47	aws.Retryer
48	Max int
49}
50
51func (w *withMaxAttempts) MaxAttempts() int {
52	return w.Max
53}
54
55// AddWithMaxBackoffDelay returns a retryer wrapping the passed in retryer
56// overriding the RetryDelay behavior for a alternate minimum initial backoff
57// delay.
58func AddWithMaxBackoffDelay(r aws.Retryer, delay time.Duration) aws.Retryer {
59	return &withMaxBackoffDelay{
60		Retryer: r,
61		backoff: NewExponentialJitterBackoff(delay),
62	}
63}
64
65type withMaxBackoffDelay struct {
66	aws.Retryer
67	backoff *ExponentialJitterBackoff
68}
69
70func (r *withMaxBackoffDelay) RetryDelay(attempt int, err error) (time.Duration, error) {
71	return r.backoff.BackoffDelay(attempt, err)
72}
73