Add support for retry mode

aws_config.go

@@ -14,6 +14,7 @@ import (
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/aws/defaults"
 	awsmiddleware "github.com/aws/aws-sdk-go-v2/aws/middleware"
 	"github.com/aws/aws-sdk-go-v2/aws/retry"
 	"github.com/aws/aws-sdk-go-v2/config"
@@ -112,23 +113,61 @@ func GetAwsConfig(ctx context.Context, c *Config) (context.Context, aws.Config,
 
 // Adapted from the per-service-client `resolveRetryer()` functions in the AWS SDK for Go v2
 // e.g. https://github.com/aws/aws-sdk-go-v2/blob/main/service/accessanalyzer/api_client.go
-// Currently only supports "standard" retry mode
 func resolveRetryer(ctx context.Context, awsConfig *aws.Config) {
-	var standardOptions []func(*retry.StandardOptions)
+	retryMode := awsConfig.RetryMode
+	if len(retryMode) == 0 {
+		defaultsMode := resolveDefaultsMode(ctx, awsConfig)
+		modeConfig, err := defaults.GetModeConfiguration(defaultsMode)
+		if err == nil {
+			retryMode = modeConfig.RetryMode
+		}
+	}
+	if len(retryMode) == 0 {
+		retryMode = aws.RetryModeStandard
+	}
 
+	var standardOptions []func(*retry.StandardOptions)
 	if v, found, _ := awsconfig.GetRetryMaxAttempts(ctx, awsConfig.ConfigSources); found && v != 0 {
 		standardOptions = append(standardOptions, func(so *retry.StandardOptions) {
 			so.MaxAttempts = v
 		})
 	}
 
+	var retryer aws.RetryerV2
+	switch retryMode {
+	case aws.RetryModeAdaptive:
+		var adaptiveOptions []func(*retry.AdaptiveModeOptions)
+		if len(standardOptions) != 0 {
+			adaptiveOptions = append(adaptiveOptions, func(ao *retry.AdaptiveModeOptions) {
+				ao.StandardOptions = append(ao.StandardOptions, standardOptions...)
+			})
+		}
+		retryer = retry.NewAdaptiveMode(adaptiveOptions...)
+
+	default:
+		retryer = retry.NewStandard(standardOptions...)
+	}
+
 	awsConfig.Retryer = func() aws.Retryer {
 		return &networkErrorShortcutter{
-			RetryerV2: retry.NewStandard(standardOptions...),
+			RetryerV2: retryer,
 		}
 	}
 }
 
+// Adapted from the per-service-client `setResolvedDefaultsMode()` functions in the AWS SDK for Go v2
+// e.g. https://github.com/aws/aws-sdk-go-v2/blob/main/service/accessanalyzer/api_client.go
+func resolveDefaultsMode(_ context.Context, awsConfig *aws.Config) aws.DefaultsMode {
+	var mode aws.DefaultsMode
+	mode.SetFromString(string(awsConfig.DefaultsMode))
+
+	if mode == aws.DefaultsModeAuto {
+		mode = defaults.ResolveDefaultsModeAuto(awsConfig.Region, awsConfig.RuntimeEnvironment)
+	}
+
+	return mode
+}
+
 // networkErrorShortcutter is used to enable networking error shortcutting
 type networkErrorShortcutter struct {
 	aws.RetryerV2
@@ -294,6 +333,12 @@ func commonLoadOptions(ctx context.Context, c *Config) ([]func(*config.LoadOptio
 		)
 	}
 
+	if c.RetryMode != "" {
+		loadOptions = append(loadOptions,
+			config.WithRetryMode(c.RetryMode),
+		)
+	}
+
 	if c.EC2MetadataServiceEndpointMode != "" {
 		var endpointMode imds.EndpointModeState
 		err := endpointMode.SetFromString(c.EC2MetadataServiceEndpointMode)

aws_config_test.go

@@ -1544,6 +1544,152 @@ max_attempts = 10
 	}
 }
 
+func TestRetryMode(t *testing.T) {
+	var (
+		standardRetryer = reflect.TypeOf((*retry.Standard)(nil))
+		adaptiveRetryer = reflect.TypeOf((*retry.AdaptiveMode)(nil))
+	)
+
+	testCases := map[string]struct {
+		Config                  *Config
+		EnvironmentVariables    map[string]string
+		SharedConfigurationFile string
+		ExpectedRetryMode       aws.RetryMode
+		RetyerType              reflect.Type
+	}{
+		"no configuration": {
+			Config: &Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+			},
+			ExpectedRetryMode: "",
+			RetyerType:        standardRetryer,
+		},
+
+		"config": {
+			Config: &Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+				RetryMode: aws.RetryModeAdaptive,
+			},
+			ExpectedRetryMode: aws.RetryModeAdaptive,
+			RetyerType:        adaptiveRetryer,
+		},
+
+		"AWS_RETRY_MODE": {
+			Config: &Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+			},
+			EnvironmentVariables: map[string]string{
+				"AWS_RETRY_MODE": "adaptive",
+			},
+			ExpectedRetryMode: aws.RetryModeAdaptive,
+			RetyerType:        adaptiveRetryer,
+		},
+
+		"shared configuration file": {
+			Config: &Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+			},
+			SharedConfigurationFile: `
+[default]
+retry_mode = adaptive
+`,
+			ExpectedRetryMode: aws.RetryModeAdaptive,
+			RetyerType:        adaptiveRetryer,
+		},
+
+		"config overrides AWS_RETRY_MODE": {
+			Config: &Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+				RetryMode: aws.RetryModeStandard,
+			},
+			EnvironmentVariables: map[string]string{
+				"AWS_RETRY_MODE": "adaptive",
+			},
+			ExpectedRetryMode: aws.RetryModeStandard,
+			RetyerType:        standardRetryer,
+		},
+
+		"AWS_RETRY_MODE overrides shared configuration": {
+			Config: &Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+			},
+			EnvironmentVariables: map[string]string{
+				"AWS_RETRY_MODE": "standard",
+			},
+			SharedConfigurationFile: `
+[default]
+retry_mode = adaptive
+`,
+			ExpectedRetryMode: aws.RetryModeStandard,
+			RetyerType:        standardRetryer,
+		},
+	}
+
+	for testName, testCase := range testCases {
+		testCase := testCase
+
+		t.Run(testName, func(t *testing.T) {
+			oldEnv := servicemocks.InitSessionTestEnv()
+			defer servicemocks.PopEnv(oldEnv)
+
+			for k, v := range testCase.EnvironmentVariables {
+				os.Setenv(k, v)
+			}
+
+			if testCase.SharedConfigurationFile != "" {
+				file, err := os.CreateTemp("", "aws-sdk-go-base-shared-configuration-file")
+
+				if err != nil {
+					t.Fatalf("unexpected error creating temporary shared configuration file: %s", err)
+				}
+
+				defer os.Remove(file.Name())
+
+				err = os.WriteFile(file.Name(), []byte(testCase.SharedConfigurationFile), 0600)
+
+				if err != nil {
+					t.Fatalf("unexpected error writing shared configuration file: %s", err)
+				}
+
+				testCase.Config.SharedConfigFiles = []string{file.Name()}
+			}
+
+			testCase.Config.SkipCredsValidation = true
+
+			_, awsConfig, err := GetAwsConfig(context.Background(), testCase.Config)
+			if err != nil {
+				t.Fatalf("error in GetAwsConfig() '%[1]T': %[1]s", err)
+			}
+
+			retryMode := awsConfig.RetryMode
+			if a, e := retryMode, testCase.ExpectedRetryMode; a != e {
+				t.Errorf(`expected RetryMode "%s", got: "%s"`, e.String(), a.String())
+			}
+
+			retryer := awsConfig.Retryer()
+			if retryer == nil {
+				t.Fatal("no retryer set")
+			}
+
+			nes, ok := retryer.(*networkErrorShortcutter)
+			if !ok {
+				t.Fatalf(`expected type "*networkErrorShortcutter", got "%T"`, retryer)
+			}
+
+			retryer = nes.RetryerV2
+			if a, e := reflect.TypeOf(retryer), testCase.RetyerType; a != e {
+				t.Errorf(`expected type "%s", got: "%s"`, e, a)
+			}
+		})
+	}
+}
+
 func TestServiceEndpointTypes(t *testing.T) {
 	testCases := map[string]struct {
 		Config                            *Config

internal/config/config.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
 	awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
 	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	"github.com/hashicorp/aws-sdk-go-base/v2/internal/expand"
@@ -35,6 +36,7 @@ type Config struct {
 	MaxRetries                     int
 	Profile                        string
 	Region                         string
+	RetryMode                      aws.RetryMode
 	SecretKey                      string
 	SharedCredentialsFiles         []string
 	SharedConfigFiles              []string

v2/awsv1shim/session_test.go

@@ -17,6 +17,7 @@ import (
 	"testing"
 	"time"
 
+	retryModev2 "github.com/aws/aws-sdk-go-v2/aws"
 	retryv2 "github.com/aws/aws-sdk-go-v2/aws/retry"
 	configv2 "github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
@@ -1408,6 +1409,125 @@ max_attempts = 10
 	}
 }
 
+func TestRetryMode(t *testing.T) {
+	testCases := map[string]struct {
+		Config                  *awsbase.Config
+		EnvironmentVariables    map[string]string
+		SharedConfigurationFile string
+		ExpectedRetryMode       retryModev2.RetryMode
+	}{
+		"no configuration": {
+			Config: &awsbase.Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+			},
+			ExpectedRetryMode: "",
+		},
+
+		"config": {
+			Config: &awsbase.Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+				RetryMode: retryModev2.RetryModeStandard,
+			},
+			ExpectedRetryMode: retryModev2.RetryModeStandard,
+		},
+
+		"AWS_RETRY_MODE": {
+			Config: &awsbase.Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+			},
+			EnvironmentVariables: map[string]string{
+				"AWS_RETRY_MODE": "adaptive",
+			},
+			ExpectedRetryMode: retryModev2.RetryModeAdaptive,
+		},
+
+		"shared configuration file": {
+			Config: &awsbase.Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+			},
+			SharedConfigurationFile: `
+		[default]
+		retry_mode = standard
+		`,
+			ExpectedRetryMode: retryModev2.RetryModeStandard,
+		},
+
+		"config overrides AWS_RETRY_MODE": {
+			Config: &awsbase.Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+				RetryMode: retryModev2.RetryModeStandard,
+			},
+			EnvironmentVariables: map[string]string{
+				"AWS_RETRY_MODE": "adaptive",
+			},
+			ExpectedRetryMode: retryModev2.RetryModeStandard,
+		},
+
+		"AWS_RETRY_MODE overrides shared configuration": {
+			Config: &awsbase.Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				SecretKey: servicemocks.MockStaticSecretKey,
+			},
+			EnvironmentVariables: map[string]string{
+				"AWS_RETRY_MODE": "standard",
+			},
+			SharedConfigurationFile: `
+		[default]
+		retry_mode = adaptive
+		`,
+			ExpectedRetryMode: retryModev2.RetryModeStandard,
+		},
+	}
+
+	for testName, testCase := range testCases {
+		testCase := testCase
+
+		t.Run(testName, func(t *testing.T) {
+			oldEnv := servicemocks.InitSessionTestEnv()
+			defer servicemocks.PopEnv(oldEnv)
+
+			for k, v := range testCase.EnvironmentVariables {
+				os.Setenv(k, v)
+			}
+
+			if testCase.SharedConfigurationFile != "" {
+				file, err := os.CreateTemp("", "aws-sdk-go-base-shared-configuration-file")
+
+				if err != nil {
+					t.Fatalf("unexpected error creating temporary shared configuration file: %s", err)
+				}
+
+				defer os.Remove(file.Name())
+
+				err = os.WriteFile(file.Name(), []byte(testCase.SharedConfigurationFile), 0600)
+
+				if err != nil {
+					t.Fatalf("unexpected error writing shared configuration file: %s", err)
+				}
+
+				testCase.Config.SharedConfigFiles = []string{file.Name()}
+			}
+
+			testCase.Config.SkipCredsValidation = true
+
+			_, awsConfig, err := awsbase.GetAwsConfig(context.Background(), testCase.Config)
+			if err != nil {
+				t.Fatalf("error in GetAwsConfig() '%[1]T': %[1]s", err)
+			}
+
+			retryMode := awsConfig.RetryMode
+			if a, e := retryMode, testCase.ExpectedRetryMode; a != e {
+				t.Errorf(`expected RetryMode "%s", got: "%s"`, e.String(), a.String())
+			}
+		})
+	}
+}
+
 func TestServiceEndpointTypes(t *testing.T) {
 	testCases := map[string]struct {
 		Config                       *awsbase.Config