go_aws_sdk_custom_retryery.go

package retryer

import (
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/client"
	"github.com/aws/aws-sdk-go/aws/request"
)

// Retryer implements a custom DefaultRetryer
type Retryer struct {
	*client.DefaultRetryer

	retryableErrorCodes []string
}

// NewRetryer constructs a DefaultRetryer with custom options
func NewRetryer(opts ...func(*Retryer)) *Retryer {
	r := &Retryer{
		DefaultRetryer: &client.DefaultRetryer{
			NumMaxRetries:    client.DefaultRetryerMaxNumRetries,
			MinRetryDelay:    client.DefaultRetryerMinRetryDelay,
			MinThrottleDelay: client.DefaultRetryerMinThrottleDelay,
			MaxRetryDelay:    client.DefaultRetryerMaxRetryDelay,
			MaxThrottleDelay: client.DefaultRetryerMaxThrottleDelay,
		},
	}

	for _, opt := range opts {
		opt(r)
	}

	return r
}

// WithMaxRetries option sets the number of time to retry requests
func WithMaxRetries(max int) func(*Retryer) {
	return func(r *Retryer) {
		r.NumMaxRetries = max
	}
}

// WithRetryableErrorCodes option adds extra error codes that aren't retryable by DefaultRetryer
func WithRetryableErrorCodes(codes ...string) func(*Retryer) {
	return func(r *Retryer) {
		r.retryableErrorCodes = append(r.retryableErrorCodes, codes...)
	}
}

// ShouldRetry is a wrapped implementation of DefaultRetryer.ShouldRetry()
func (r *Retryer) ShouldRetry(req *request.Request) bool {
	if aerr, ok := req.Error.(awserr.Error); ok && aerr != nil {
		for _, code := range r.retryableErrorCodes {
			if code == aerr.Code() {
				req.Retryable = aws.Bool(true)
			}
		}
	}

	return r.DefaultRetryer.ShouldRetry(req)
}