Skip to content

Commit 814084f

Browse files
mbambermbamber
committed
feat: Allow customizing of AWS retry codes
Allow users to customize the error codes that should be retried by the AWS SDK. This enables advanced workflows such as retrying authentication failures
1 parent 5bc808b commit 814084f

File tree

4 files changed

+200
-0
lines changed

4 files changed

+200
-0
lines changed

‎aws_config.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@ func GetAwsConfig(ctx context.Context, c *Config) (aws.Config, error) {
5252
if c.MaxRetries != 0 {
5353
retryer = retry.AddWithMaxAttempts(retryer, c.MaxRetries)
5454
}
55+
if retryCodes := os.Getenv("AWS_RETRY_CODES"); retryCodes != "" {
56+
codes := strings.Split(retryCodes, ",")
57+
retryer = retry.AddWithErrorCodes(retryer, codes...)
58+
}
5559
retryer = &networkErrorShortcutter{
5660
Retryer: retryer,
5761
}

‎aws_config_test.go‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"github.com/aws/aws-sdk-go-v2/config"
2121
"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
2222
"github.com/aws/aws-sdk-go-v2/service/sts"
23+
"github.com/aws/smithy-go"
2324
"github.com/aws/smithy-go/middleware"
2425
smithyhttp "github.com/aws/smithy-go/transport/http"
2526
"github.com/google/go-cmp/cmp"
@@ -1630,6 +1631,93 @@ func TestMaxAttempts(t *testing.T) {
16301631
}
16311632
}
16321633

1634+
func TestRetryCodes(t *testing.T) {
1635+
testCases := map[string]struct {
1636+
Config *Config
1637+
EnvironmentVariables map[string]string
1638+
ExpectedRetryableErrors []smithy.APIError
1639+
ExpectedNonRetryableErrors []smithy.APIError
1640+
}{
1641+
"no configuration": {
1642+
Config: &Config{
1643+
AccessKey: servicemocks.MockStaticAccessKey,
1644+
SecretKey: servicemocks.MockStaticSecretKey,
1645+
},
1646+
ExpectedNonRetryableErrors: []smithy.APIError{
1647+
&smithy.GenericAPIError{Code: "error 1"},
1648+
},
1649+
},
1650+
1651+
"AWS_RETRY_CODES single": {
1652+
Config: &Config{
1653+
AccessKey: servicemocks.MockStaticAccessKey,
1654+
SecretKey: servicemocks.MockStaticSecretKey,
1655+
},
1656+
EnvironmentVariables: map[string]string{
1657+
"AWS_RETRY_CODES": "error 1",
1658+
},
1659+
ExpectedRetryableErrors: []smithy.APIError{
1660+
&smithy.GenericAPIError{Code: "error 1"},
1661+
},
1662+
ExpectedNonRetryableErrors: []smithy.APIError{
1663+
&smithy.GenericAPIError{Code: "error 2"},
1664+
},
1665+
},
1666+
1667+
"AWS_RETRY_CODES multiple": {
1668+
Config: &Config{
1669+
AccessKey: servicemocks.MockStaticAccessKey,
1670+
SecretKey: servicemocks.MockStaticSecretKey,
1671+
},
1672+
EnvironmentVariables: map[string]string{
1673+
"AWS_RETRY_CODES": "error 1,error 2",
1674+
},
1675+
ExpectedRetryableErrors: []smithy.APIError{
1676+
&smithy.GenericAPIError{Code: "error 1"},
1677+
&smithy.GenericAPIError{Code: "error 2"},
1678+
},
1679+
ExpectedNonRetryableErrors: []smithy.APIError{
1680+
&smithy.GenericAPIError{Code: "error 3"},
1681+
},
1682+
},
1683+
}
1684+
1685+
for testName, testCase := range testCases {
1686+
testCase := testCase
1687+
1688+
t.Run(testName, func(t *testing.T) {
1689+
oldEnv := servicemocks.InitSessionTestEnv()
1690+
defer servicemocks.PopEnv(oldEnv)
1691+
1692+
for k, v := range testCase.EnvironmentVariables {
1693+
os.Setenv(k, v)
1694+
}
1695+
1696+
testCase.Config.SkipCredsValidation = true
1697+
1698+
awsConfig, err := GetAwsConfig(context.Background(), testCase.Config)
1699+
if err != nil {
1700+
t.Fatalf("error in GetAwsConfig() '%[1]T': %[1]s", err)
1701+
}
1702+
1703+
retryer := awsConfig.Retryer()
1704+
if retryer == nil {
1705+
t.Fatal("no retryer set")
1706+
}
1707+
for _, e := range testCase.ExpectedRetryableErrors {
1708+
if a := retryer.IsErrorRetryable(e); !a {
1709+
t.Errorf(`expected error %q would be retryable, got not retryable`, e)
1710+
}
1711+
}
1712+
for _, e := range testCase.ExpectedNonRetryableErrors {
1713+
if a := retryer.IsErrorRetryable(e); a {
1714+
t.Errorf(`expected error %q would not be retryable, got retryable`, e)
1715+
}
1716+
}
1717+
})
1718+
}
1719+
}
1720+
16331721
func TestServiceEndpointTypes(t *testing.T) {
16341722
testCases := map[string]struct {
16351723
Config *Config

‎v2/awsv1shim/session.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import ( // nosemgrep: no-sdkv2-imports-in-awsv1shim
55
"fmt"
66
"log"
77
"os"
8+
"strings"
89

910
awsv2 "github.com/aws/aws-sdk-go-v2/aws"
1011
"github.com/aws/aws-sdk-go/aws"
@@ -89,6 +90,18 @@ func GetSession(awsC *awsv2.Config, c *awsbase.Config) (*session.Session, error)
8990
sess = sess.Copy(&aws.Config{MaxRetries: aws.Int(retryer.MaxAttempts())})
9091
}
9192

93+
// Add custom error code retries. It's easier to recheck the environment variable
94+
// here as the retry codes aren't available from the original v2 config
95+
if retryCodes := os.Getenv("AWS_RETRY_CODES"); retryCodes != "" {
96+
codes := strings.Split(retryCodes, ",")
97+
log.Printf("[DEBUG] Using additional retry codes: %s", codes)
98+
sess.Handlers.Retry.PushBack(func(r *request.Request) {
99+
if tfawserr.ErrCodeEquals(r.Error, codes...) {
100+
r.Retryable = aws.Bool(true)
101+
}
102+
})
103+
}
104+
92105
SetSessionUserAgent(sess, c.APNInfo, c.UserAgent)
93106

94107
// Add custom input from ENV to the User-Agent request header

‎v2/awsv1shim/session_test.go‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1450,6 +1450,101 @@ func TestMaxAttempts(t *testing.T) {
14501450
}
14511451
}
14521452

1453+
func TestRetryCodes(t *testing.T) {
1454+
testCases := map[string]struct {
1455+
Config *awsbase.Config
1456+
EnvironmentVariables map[string]string
1457+
ExpectedRetryableErrors []awserr.Error
1458+
ExpectedNonRetryableErrors []awserr.Error
1459+
}{
1460+
"no configuration": {
1461+
Config: &awsbase.Config{
1462+
AccessKey: servicemocks.MockStaticAccessKey,
1463+
SecretKey: servicemocks.MockStaticSecretKey,
1464+
},
1465+
ExpectedNonRetryableErrors: []awserr.Error{
1466+
awserr.New("error 1", "", nil),
1467+
},
1468+
},
1469+
1470+
"AWS_RETRY_CODES single": {
1471+
Config: &awsbase.Config{
1472+
AccessKey: servicemocks.MockStaticAccessKey,
1473+
SecretKey: servicemocks.MockStaticSecretKey,
1474+
},
1475+
EnvironmentVariables: map[string]string{
1476+
"AWS_RETRY_CODES": "error 1",
1477+
},
1478+
ExpectedRetryableErrors: []awserr.Error{
1479+
awserr.New("error 1", "", nil),
1480+
},
1481+
ExpectedNonRetryableErrors: []awserr.Error{
1482+
awserr.New("error 2", "", nil),
1483+
},
1484+
},
1485+
1486+
"AWS_RETRY_CODES multiple": {
1487+
Config: &awsbase.Config{
1488+
AccessKey: servicemocks.MockStaticAccessKey,
1489+
SecretKey: servicemocks.MockStaticSecretKey,
1490+
},
1491+
EnvironmentVariables: map[string]string{
1492+
"AWS_RETRY_CODES": "error 1,error 2",
1493+
},
1494+
ExpectedRetryableErrors: []awserr.Error{
1495+
awserr.New("error 1", "", nil),
1496+
awserr.New("error 2", "", nil),
1497+
},
1498+
ExpectedNonRetryableErrors: []awserr.Error{
1499+
awserr.New("error 3", "", nil),
1500+
},
1501+
},
1502+
}
1503+
1504+
for testName, testCase := range testCases {
1505+
testCase := testCase
1506+
1507+
t.Run(testName, func(t *testing.T) {
1508+
oldEnv := servicemocks.InitSessionTestEnv()
1509+
defer servicemocks.PopEnv(oldEnv)
1510+
1511+
for k, v := range testCase.EnvironmentVariables {
1512+
os.Setenv(k, v)
1513+
}
1514+
1515+
testCase.Config.SkipCredsValidation = true
1516+
1517+
awsConfig, err := awsbase.GetAwsConfig(context.Background(), testCase.Config)
1518+
if err != nil {
1519+
t.Fatalf("GetAwsConfig() returned error: %s", err)
1520+
}
1521+
actualSession, err := GetSession(&awsConfig, testCase.Config)
1522+
if err != nil {
1523+
t.Fatalf("error in GetSession() '%[1]T': %[1]s", err)
1524+
}
1525+
1526+
for _, e := range testCase.ExpectedRetryableErrors {
1527+
r := &request.Request{
1528+
Error: e,
1529+
}
1530+
actualSession.Handlers.Retry.Run(r)
1531+
if !aws.BoolValue(r.Retryable) {
1532+
t.Errorf(`expected error %q would be retryable, got not retryable`, e)
1533+
}
1534+
}
1535+
for _, e := range testCase.ExpectedNonRetryableErrors {
1536+
r := &request.Request{
1537+
Error: e,
1538+
}
1539+
actualSession.Handlers.Retry.Run(r)
1540+
if aws.BoolValue(r.Retryable) {
1541+
t.Errorf(`expected error %q would not be retryable, got retryable`, e)
1542+
}
1543+
}
1544+
})
1545+
}
1546+
}
1547+
14531548
func TestServiceEndpointTypes(t *testing.T) {
14541549
testCases := map[string]struct {
14551550
Config *awsbase.Config

0 commit comments

Comments
 (0)