Add preview of changes for standard retry mode behind flag (#3400)

* Add changes for new retries 2026 behind a flag

* Add changelog

* Make new jitter options private and instead expose BaseDelay for DDB use case; Remove LongPoll from retry options

* Add STS IDPCommunicationError as a retryable error

* Update long polling codegen

* Add new test for exhausted quota on long-polling operations

Author: Madrigal

.changelog/cffa1fca582844fb947bae754cd08bae.json

@@ -0,0 +1,13 @@
+{
+    "id": "cffa1fca-5828-44fb-947b-ae754cd08bae",
+    "type": "feature",
+    "description": "Add preview of standard retry changes behind AWS_NEW_RETRIES_2026 flag",
+    "modules": [
+        ".",
+        "service/dynamodb",
+        "service/dynamodbstreams",
+        "service/sfn",
+        "service/sqs",
+        "service/swf"
+    ]
+}

aws/retry/jitter_backoff.go

@@ -4,6 +4,7 @@ import (
 	"math"
 	"time"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/internal/rand"
 	"github.com/aws/aws-sdk-go-v2/internal/timeconv"
 )
@@ -12,9 +13,20 @@ import (
 // number of attempts.
 type ExponentialJitterBackoff struct {
 	maxBackoff time.Duration
-	// precomputed number of attempts needed to reach max backoff.
+	// precomputed number of attempts needed to reach max backoff (legacy mode).
 	maxBackoffAttempts float64
 
+	// Base delay for non-throttle errors (x in the formula t_i = b * min(x * r^i, MAX_BACKOFF)).
+	baseDelay time.Duration
+
+	// Throttle error checker. When set and the error is a throttle, the base
+	// delay is 1s regardless of the configured baseDelay.
+	throttle IsErrorThrottle
+
+	// When true, applies MAX_BACKOFF before jitter and uses throttle-aware
+	// base delay.
+	retries2026 bool
+
 	randFloat64 func() (float64, error)
 }
 
@@ -25,13 +37,53 @@ func NewExponentialJitterBackoff(maxBackoff time.Duration) *ExponentialJitterBac
 		maxBackoff: maxBackoff,
 		maxBackoffAttempts: math.Log2(
 			float64(maxBackoff) / float64(time.Second)),
+		baseDelay:   time.Second,
 		randFloat64: rand.CryptoRandFloat64,
 	}
 }
 
+// exponentialJitterBackoffOption is a functional option for ExponentialJitterBackoff.
+type exponentialJitterBackoffOption func(*ExponentialJitterBackoff)
+
+// withBaseDelay sets the base delay for non-throttle errors.
+func withBaseDelay(d time.Duration) exponentialJitterBackoffOption {
+	return func(j *ExponentialJitterBackoff) {
+		j.baseDelay = d
+	}
+}
+
+// withThrottleCheck sets the throttle error checker used to determine if the
+// backoff should use the throttle base delay (1s) instead of the configured
+// base delay.
+func withThrottleCheck(t IsErrorThrottle) exponentialJitterBackoffOption {
+	return func(j *ExponentialJitterBackoff) {
+		j.throttle = t
+	}
+}
+
+// newExponentialJitterBackoffWithOptions returns an ExponentialJitterBackoff
+// with the given options applied.
+func newExponentialJitterBackoffWithOptions(maxBackoff time.Duration, optFns ...exponentialJitterBackoffOption) *ExponentialJitterBackoff {
+	j := NewExponentialJitterBackoff(maxBackoff)
+	j.retries2026 = true
+	for _, fn := range optFns {
+		fn(j)
+	}
+	return j
+}
+
 // BackoffDelay returns the duration to wait before the next attempt should be
 // made. Returns an error if unable get a duration.
 func (j *ExponentialJitterBackoff) BackoffDelay(attempt int, err error) (time.Duration, error) {
+	if j.retries2026 {
+		return j.backoffDelay2026(attempt, err)
+	}
+	return j.backoffDelayLegacy(attempt, err)
+}
+
+// backoffDelayLegacy preserves the original backoff formula: b * 2^i, capped
+// at maxBackoff.
+func (j *ExponentialJitterBackoff) backoffDelayLegacy(attempt int, err error) (time.Duration, error) {
 	if attempt > int(j.maxBackoffAttempts) {
 		return j.maxBackoff, nil
 	}
@@ -47,3 +99,26 @@ func (j *ExponentialJitterBackoff) BackoffDelay(attempt int, err error) (time.Du
 
 	return timeconv.FloatSecondsDur(delaySeconds), nil
 }
+
+// backoffDelay2026 uses throttle-aware base delay and applies MAX_BACKOFF
+// before jitter: t_i = b * min(x * 2^i, MAX_BACKOFF).
+func (j *ExponentialJitterBackoff) backoffDelay2026(attempt int, err error) (time.Duration, error) {
+	x := j.baseDelay
+	if j.throttle != nil && j.throttle.IsErrorThrottle(err) == aws.TrueTernary {
+		x = time.Second
+	}
+
+	b, randErr := j.randFloat64()
+	if randErr != nil {
+		return 0, randErr
+	}
+
+	ri := math.Pow(2, float64(attempt))
+	delaySeconds := float64(x) / float64(time.Second) * ri
+	maxBackoffSeconds := float64(j.maxBackoff) / float64(time.Second)
+	if delaySeconds > maxBackoffSeconds {
+		delaySeconds = maxBackoffSeconds
+	}
+
+	return timeconv.FloatSecondsDur(b * delaySeconds), nil
+}

aws/retry/middleware.go

@@ -233,9 +233,11 @@ func (r *Attempt) handleAttempt(
 			"failed to release retry token after request error, %w", err)
 	}
 	// Release the attempt token based on the state of the attempt's error (if any).
-	if releaseError := releaseAttemptToken(err); releaseError != nil && err != nil {
-		return out, attemptResult, nopRelease, fmt.Errorf(
-			"failed to release initial token after request error, %w", err)
+	if !newRetries2026() || attemptNum == 1 {
+		if releaseError := releaseAttemptToken(err); releaseError != nil && err != nil {
+			return out, attemptResult, nopRelease, fmt.Errorf(
+				"failed to release initial token after request error, %w", err)
+		}
 	}
 	// If there was no error making the attempt, nothing further to do. There
 	// will be nothing to retry.
@@ -276,6 +278,13 @@ func (r *Attempt) handleAttempt(
 	// Get a retry token that will be released after the
 	releaseRetryToken, retryTokenErr := r.retryer.GetRetryToken(ctx, err)
 	if retryTokenErr != nil {
+		// Long-polling operations must still back off when quota is exceeded.
+		if newRetries2026() && internalcontext.GetIsLongPolling(ctx) {
+			if retryDelay, delayErr := r.retryer.RetryDelay(attemptNum-1, err); delayErr == nil {
+				retryDelay = adjustForRetryAfterHeader(retryDelay, err, logger, r.LogAttempts)
+				_ = sdk.SleepWithContext(ctx, retryDelay)
+			}
+		}
 		return out, attemptResult, nopRelease, errors.Join(err, retryTokenErr)
 	}
 
@@ -285,10 +294,17 @@ func (r *Attempt) handleAttempt(
 	// Get the retry delay before another attempt can be made, and sleep for
 	// that time. Potentially early exist if the sleep is canceled via the
 	// context.
-	retryDelay, reqErr := r.retryer.RetryDelay(attemptNum, err)
+	attempt := attemptNum
+	if newRetries2026() {
+		attempt = attemptNum - 1
+	}
+	retryDelay, reqErr := r.retryer.RetryDelay(attempt, err)
 	if reqErr != nil {
 		return out, attemptResult, releaseRetryToken, reqErr
 	}
+	if newRetries2026() {
+		retryDelay = adjustForRetryAfterHeader(retryDelay, err, logger, r.LogAttempts)
+	}
 	if reqErr = sdk.SleepWithContext(ctx, retryDelay); reqErr != nil {
 		err = &aws.RequestCanceledError{Err: reqErr}
 		return out, attemptResult, releaseRetryToken, err
@@ -423,6 +439,43 @@ func AddRetryMiddlewares(stack *smithymiddle.Stack, options AddRetryMiddlewaresO
 	return nil
 }
 
+// adjustForRetryAfterHeader checks for the x-amz-retry-after response header
+// and clamps the backoff duration accordingly. The header value is an integer
+// representing milliseconds. The result is clamped to [t_i, 5s + t_i] where
+// t_i is the jittered exponential backoff duration. Invalid header values are
+// ignored.
+func adjustForRetryAfterHeader(backoff time.Duration, err error, logger logging.Logger, logAttempts bool) time.Duration {
+	var re *http.ResponseError
+	if !errors.As(err, &re) || re.Response == nil || re.Response.Response == nil {
+		return backoff
+	}
+
+	headerVal := re.Response.Header.Get("X-Amz-Retry-After")
+	if headerVal == "" {
+		return backoff
+	}
+
+	ms, parseErr := strconv.ParseInt(headerVal, 10, 64)
+	if parseErr != nil || ms < 0 {
+		if logAttempts {
+			logger.Logf(logging.Debug, "ignoring invalid x-amz-retry-after header value %q", headerVal)
+		}
+		return backoff
+	}
+
+	retryAfter := time.Duration(ms) * time.Millisecond
+	minDuration := backoff
+	maxDuration := 5*time.Second + backoff
+
+	if retryAfter < minDuration {
+		return minDuration
+	}
+	if retryAfter > maxDuration {
+		return maxDuration
+	}
+	return retryAfter
+}
+
 // Determines the value of exception.type for metrics purposes. We prefer an
 // API-specific error code, otherwise it's just the Go type for the value.
 func errorType(err error) string {

aws/retry/middleware_test.go

@@ -564,6 +564,58 @@ func TestClockSkew(t *testing.T) {
 	}
 }
 
+// TestLegacyRetryDelayAttemptValue verifies that when the new retries flag is
+// NOT set, the attempt value passed to RetryDelay is unchanged from legacy
+// behavior (1-based: first retry passes 1, second passes 2, etc.).
+func TestLegacyRetryDelayAttemptValue(t *testing.T) {
+	restoreSleep := sdk.TestingUseNopSleep()
+	defer restoreSleep()
+
+	var recordedAttempts []int
+	retryer := NewStandard(func(o *StandardOptions) {
+		o.MaxAttempts = 4
+		o.Backoff = BackoffDelayerFunc(func(attempt int, err error) (time.Duration, error) {
+			recordedAttempts = append(recordedAttempts, attempt)
+			return 0, nil
+		})
+	})
+
+	am := NewAttemptMiddleware(retryer, func(i interface{}) interface{} {
+		return i
+	})
+
+	num := 0
+	handler := middleware.FinalizeHandlerFunc(
+		func(ctx context.Context, in middleware.FinalizeInput) (
+			out middleware.FinalizeOutput, metadata middleware.Metadata, err error,
+		) {
+			num++
+			if num < 4 {
+				return out, metadata, mockRetryableError{b: true}
+			}
+			return out, metadata, nil
+		})
+
+	_, _, err := am.HandleFinalize(context.Background(), middleware.FinalizeInput{}, handler)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	// Legacy behavior: attemptNum is passed directly (1-based).
+	// After attempt 1 fails, RetryDelay is called with attemptNum=1.
+	// After attempt 2 fails, RetryDelay is called with attemptNum=2.
+	// After attempt 3 fails, RetryDelay is called with attemptNum=3.
+	expected := []int{1, 2, 3}
+	if len(recordedAttempts) != len(expected) {
+		t.Fatalf("expected %d RetryDelay calls, got %d", len(expected), len(recordedAttempts))
+	}
+	for i, exp := range expected {
+		if recordedAttempts[i] != exp {
+			t.Errorf("RetryDelay call %d: expected attempt=%d, got attempt=%d", i, exp, recordedAttempts[i])
+		}
+	}
+}
+
 // mockRawResponseKey is used to test the behavior when response metadata is
 // nested within the attempt request.
 type mockRawResponseKey struct{}

aws/retry/retry.go

@@ -72,6 +72,19 @@ func (r *withMaxBackoffDelay) RetryDelay(attempt int, err error) (time.Duration,
 	return r.backoff.BackoffDelay(attempt, err)
 }
 
+// AddWithLongPolling returns a retryer that is marked as long-polling.
+// Long-polling operations will back off even when the retry quota is
+// exhausted.
+func AddWithLongPolling(r aws.Retryer) aws.Retryer {
+	return &withLongPolling{RetryerV2: wrapAsRetryerV2(r)}
+}
+
+type withLongPolling struct {
+	aws.RetryerV2
+}
+
+func (w *withLongPolling) IsLongPolling() bool { return true }
+
 type wrappedAsRetryerV2 struct {
 	aws.Retryer
 }