Implement updated standard retry behavior

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsDynamoDbRetryIntegration.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.List;
+import java.util.Set;
+import software.amazon.smithy.aws.traits.ServiceTrait;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.RuntimeTypes;
+import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
+import software.amazon.smithy.python.codegen.sections.InitRetryStrategyResolverSection;
+import software.amazon.smithy.python.codegen.writer.PythonWriter;
+import software.amazon.smithy.utils.CodeInterceptor;
+import software.amazon.smithy.utils.CodeSection;
+
+/**
+ * Injects DynamoDB's default retry options (max attempts 4, 25ms non-throttling
+ * base backoff).
+ */
+public final class AwsDynamoDbRetryIntegration implements PythonIntegration {
+
+    private static final Set<String> DYNAMODB_SDK_IDS = Set.of("DynamoDB", "DynamoDB Streams");
+
+    private static final double DYNAMODB_BASE_BACKOFF_SECONDS = 0.025;
+    private static final int DYNAMODB_MAX_ATTEMPTS = 4;
+
+    private static boolean isDynamoDb(GenerationContext context) {
+        return context.settings()
+                .service(context.model())
+                .getTrait(ServiceTrait.class)
+                .map(trait -> DYNAMODB_SDK_IDS.contains(trait.getSdkId()))
+                .orElse(false);
+    }
+
+    @Override
+    public List<? extends CodeInterceptor<? extends CodeSection, PythonWriter>> interceptors(
+            GenerationContext context
+    ) {
+        if (!isDynamoDb(context)) {
+            return List.of();
+        }
+        return List.of(new DynamoDbRetryStrategyResolverInterceptor());
+    }
+
+    private static final class DynamoDbRetryStrategyResolverInterceptor
+            implements CodeInterceptor<InitRetryStrategyResolverSection, PythonWriter> {
+
+        @Override
+        public Class<InitRetryStrategyResolverSection> sectionType() {
+            return InitRetryStrategyResolverSection.class;
+        }
+
+        @Override
+        public void write(PythonWriter writer, String previousText, InitRetryStrategyResolverSection section) {
+            writer.write(
+                    "self._retry_strategy_resolver = $T(default_max_attempts=$L, default_backoff_scale=$L)",
+                    RuntimeTypes.RETRY_STRATEGY_RESOLVER,
+                    DYNAMODB_MAX_ATTEMPTS,
+                    DYNAMODB_BASE_BACKOFF_SECONDS);
+        }
+    }
+}

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsLongPollingIntegration.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.Map;
+import java.util.Set;
+import software.amazon.smithy.aws.traits.ServiceTrait;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.traits.LongPollTrait;
+import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
+
+/**
+ * Marks long-polling operations so the generic client generator applies
+ * long-polling retry behavior to them.
+ *
+ * <p>An operation is long-polling if it carries the {@code smithy.api#longPoll}
+ * trait. Service models do not yet apply the trait, so the known operations are
+ * hard-coded as a fallback. Once the trait ships in the models, the fallback
+ * can be removed.
+ */
+public final class AwsLongPollingIntegration implements PythonIntegration {
+
+    private static final Map<String, Set<String>> LONG_POLLING_OPERATIONS = Map.of(
+            "SQS",
+            Set.of("ReceiveMessage"),
+            "SFN",
+            Set.of("GetActivityTask"),
+            "SWF",
+            Set.of("PollForActivityTask", "PollForDecisionTask"));
+
+    @Override
+    public boolean isLongPollingOperation(Model model, ServiceShape service, OperationShape operation) {
+        if (operation.hasTrait(LongPollTrait.class)) {
+            return true;
+        }
+        return service.getTrait(ServiceTrait.class)
+                .map(trait -> LONG_POLLING_OPERATIONS.get(trait.getSdkId()))
+                .map(operations -> operations.contains(operation.getId().getName()))
+                .orElse(false);
+    }
+}

codegen/aws/core/src/main/resources/META-INF/services/software.amazon.smithy.python.codegen.integrations.PythonIntegration

@@ -8,3 +8,5 @@ software.amazon.smithy.python.aws.codegen.AwsProtocolsIntegration
 software.amazon.smithy.python.aws.codegen.AwsServiceIdIntegration
 software.amazon.smithy.python.aws.codegen.AwsUserAgentIntegration
 software.amazon.smithy.python.aws.codegen.AwsStandardRegionalEndpointsIntegration
+software.amazon.smithy.python.aws.codegen.AwsDynamoDbRetryIntegration
+software.amazon.smithy.python.aws.codegen.AwsLongPollingIntegration

codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java

@@ -22,6 +22,7 @@
 import software.amazon.smithy.model.traits.StringTrait;
 import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
 import software.amazon.smithy.python.codegen.integrations.RuntimeClientPlugin;
+import software.amazon.smithy.python.codegen.sections.InitRetryStrategyResolverSection;
 import software.amazon.smithy.python.codegen.writer.PythonWriter;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
@@ -86,13 +87,13 @@ def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None):
                         for plugin in client_plugins:
                             plugin(self._config)
 
-                        self._retry_strategy_resolver = $5T()
+                        $5C
                     """,
                     configSymbol,
                     pluginSymbol,
                     writer.consumer(w -> writeConstructorDocs(w, serviceSymbol.getName())),
                     writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)),
-                    RuntimeTypes.RETRY_STRATEGY_RESOLVER);
+                    writer.consumer(this::writeRetryStrategyResolverInit));
 
             var topDownIndex = TopDownIndex.of(model);
             var eventStreamIndex = EventStreamIndex.of(model);
@@ -113,6 +114,22 @@ private void writeDefaultPlugins(PythonWriter writer, Collection<SymbolReference
         }
     }
 
+    private void writeRetryStrategyResolverInit(PythonWriter writer) {
+        // Wrap this in a section so AWS integrations can inject service-specific retry defaults.
+        writer.pushState(new InitRetryStrategyResolverSection());
+        writer.write("self._retry_strategy_resolver = $T()", RuntimeTypes.RETRY_STRATEGY_RESOLVER);
+        writer.popState();
+    }
+
+    private boolean isLongPollingOperation(OperationShape operation) {
+        for (PythonIntegration integration : context.integrations()) {
+            if (integration.isLongPollingOperation(model, service, operation)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     private void writeConstructorDocs(PythonWriter writer, String clientName) {
         writer.writeMultiLineDocs(() -> {
             writer.write("""
@@ -210,6 +227,7 @@ private void writeSharedOperationInit(
         }
 
         writer.putContext("operation", symbolProvider.toSymbol(operation));
+        writer.putContext("isLongPolling", isLongPollingOperation(operation));
         writer.addStdlibImport("copy", "deepcopy");
 
         writer.write("""
@@ -235,8 +253,8 @@ private void writeSharedOperationInit(
                 call = $4T(
                     input=input,
                     operation=${operation:T},
-                    context=$5T({"config": config}),
-                    interceptor=$6T(config.interceptors),
+                    context=$5T({"config": config${?isLongPolling}, $6T: True${/isLongPolling}}),
+                    interceptor=$7T(config.interceptors),
                     auth_scheme_resolver=config.auth_scheme_resolver,
                     supported_auth_schemes=config.auth_schemes,
                     endpoint_resolver=config.endpoint_resolver,
@@ -248,6 +266,7 @@ private void writeSharedOperationInit(
                 RuntimeTypes.REQUEST_PIPELINE,
                 RuntimeTypes.CLIENT_CALL,
                 RuntimeTypes.TYPED_PROPERTIES,
+                RuntimeTypes.LONG_POLLING,
                 RuntimeTypes.INTERCEPTOR_CHAIN);
 
     }

codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java

@@ -68,6 +68,8 @@ public final class RuntimeTypes {
             createSymbol("aio.retries", "RetryStrategyResolver", SmithyPythonDependency.SMITHY_CORE);
     public static final Symbol RETRY_STRATEGY_OPTIONS =
             createSymbol("retries", "RetryStrategyOptions", SmithyPythonDependency.SMITHY_CORE);
+    public static final Symbol LONG_POLLING =
+            createSymbol("retries", "LONG_POLLING", SmithyPythonDependency.SMITHY_CORE);
     public static final Symbol SIMPLE_RETRY_STRATEGY =
             createSymbol("aio.retries", "SimpleRetryStrategy", SmithyPythonDependency.SMITHY_CORE);
 

codegen/core/src/main/java/software/amazon/smithy/python/codegen/integrations/PythonIntegration.java

@@ -8,6 +8,8 @@
 import java.util.List;
 import software.amazon.smithy.codegen.core.SmithyIntegration;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.python.codegen.GenerationContext;
 import software.amazon.smithy.python.codegen.PythonSettings;
 import software.amazon.smithy.python.codegen.generators.ProtocolGenerator;
@@ -44,6 +46,21 @@ default Model preprocessModel(Model model, PythonSettings settings) {
         return model;
     }
 
+    /**
+     * Determines whether the given operation is a long-polling operation, which
+     * must back off before returning even when the retry quota is exhausted. AWS
+     * integrations use this hook to identify these operations while the
+     * {@code aws.api#longPoll} trait is not yet shipped in service models.
+     *
+     * @param model     Model the operation belongs to.
+     * @param service   Service the operation belongs to.
+     * @param operation Operation to test.
+     * @return Returns true if the operation is a long-polling operation.
+     */
+    default boolean isLongPollingOperation(Model model, ServiceShape service, OperationShape operation) {
+        return false;
+    }
+
     /**
      * Writes out all extra files required by runtime plugins.
      */

codegen/core/src/main/java/software/amazon/smithy/python/codegen/sections/InitRetryStrategyResolverSection.java

@@ -0,0 +1,14 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.codegen.sections;
+
+import software.amazon.smithy.utils.CodeSection;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+/**
+ * A section that controls constructing the client's {@code RetryStrategyResolver}.
+ */
+@SmithyInternalApi
+public record InitRetryStrategyResolverSection() implements CodeSection {}

designs/retries.md

@@ -32,25 +32,33 @@ class RetryStrategy(Protocol):
     """Upper limit on total attempt count (initial attempt plus retries)."""
 
     async def acquire_initial_retry_token(
-        self, *, token_scope: str | None = None
+        self, *, token_scope: str | None = None, context: TypedProperties | None = None
     ) -> RetryToken:
         """Called before any retries (for the first attempt at the operation).
 
         :param token_scope: An arbitrary string accepted by the retry strategy to
             separate tokens into scopes.
+        :param context: The operation context, read for per-operation signals such as
+            :py:data:`smithy_core.retries.LONG_POLLING`.
         :returns: A retry token, to be used for determining the retry delay, refreshing
             the token after a failure, and recording success after success.
         :raises RetryError: If the retry strategy has no available tokens.
         """
         ...
 
     async def refresh_retry_token_for_retry(
-        self, *, token_to_renew: RetryToken, error: Exception
+        self,
+        *,
+        token_to_renew: RetryToken,
+        error: Exception,
+        context: TypedProperties | None = None,
     ) -> RetryToken:
         """Replace an existing retry token from a failed attempt with a new token.
 
         :param token_to_renew: The token used for the previous failed attempt.
         :param error: The error that triggered the need for a retry.
+        :param context: The operation context, read for per-operation signals such as
+            :py:data:`smithy_core.retries.LONG_POLLING`.
         :raises RetryError: If no further retry attempts are allowed.
         """
         ...
@@ -110,8 +118,9 @@ class HasFault(Protocol):
 
 `RetryStrategy` implementations MUST raise a `RetryError` if they receive an
 exception where `is_retry_safe` is `False` and SHOULD raise a `RetryError` if it
-is `None`. `RetryStrategy` implementations SHOULD use a delay that is at least
-as long as `retry_after` but MAY choose to wait longer.
+is `None`. `RetryStrategy` implementations SHOULD take `retry_after` into account
+when computing the delay, but MAY adjust it (for example, by clamping it to an
+upper bound).
 
 ### Backoff Strategy
 

packages/smithy-aws-core/.changes/next-release/smithy-aws-core-feature-881860b5bda049d1b000983dd2a3bd0e.json

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "Added support for the `x-amz-retry-after` response header."
+}

packages/smithy-aws-core/src/smithy_aws_core/_private/query/errors.py

@@ -99,6 +99,7 @@ def create_aws_query_error(
     wrapper_elements: tuple[str, ...],
     status: int,
     context: TypedProperties,
+    retry_after: float | None = None,
 ) -> CallError:
     """Create a modeled or generic CallError from an awsQuery error response."""
     code = _parse_aws_query_error_code(body, wrapper_elements)
@@ -121,7 +122,10 @@ def create_aws_query_error(
             deserializer = XMLCodec().create_deserializer(
                 body, wrapper_elements=wrapper_elements
             )
-            return error_shape.deserialize(deserializer)
+            modeled_error = error_shape.deserialize(deserializer)
+            if retry_after is not None:
+                modeled_error.retry_after = retry_after
+            return modeled_error
 
     message = f"Unknown error for operation {operation.schema.id} - status: {status}"
     if code is not None:
@@ -137,4 +141,5 @@ def create_aws_query_error(
         is_throttling_error=is_throttle,
         is_timeout_error=is_timeout,
         is_retry_safe=is_throttle or is_timeout or None,
+        retry_after=retry_after,
     )

packages/smithy-aws-core/src/smithy_aws_core/aio/protocols.py

@@ -27,15 +27,18 @@
 from smithy_http import tuples_to_fields
 from smithy_http.aio import HTTPRequest as _HTTPRequest
 from smithy_http.aio.interfaces import HTTPErrorIdentifier, HTTPRequest, HTTPResponse
-from smithy_http.aio.protocols import HttpBindingClientProtocol, HttpClientProtocol
+from smithy_http.aio.protocols import (
+    HttpBindingClientProtocol,
+    HttpClientProtocol,
+)
 from smithy_http.deserializers import HTTPResponseDeserializer
 
 from .._private.query.errors import (
     create_aws_query_error,
 )
 from .._private.query.serializers import QueryShapeSerializer
 from ..traits import AwsQueryTrait, RestJson1Trait
-from ..utils import parse_document_discriminator, parse_error_code
+from ..utils import parse_document_discriminator, parse_error_code, parse_retry_after
 
 try:
     from smithy_json import JSONCodec, JSONDocument
@@ -166,6 +169,9 @@ def content_type(self) -> str:
     def error_identifier(self) -> HTTPErrorIdentifier:
         return self._error_identifier
 
+    def _retry_after(self, response: HTTPResponse) -> float | None:
+        return parse_retry_after(response)
+
     def create_event_publisher[
         OperationInput: SerializeableShape,
         OperationOutput: DeserializeableShape,
@@ -353,6 +359,7 @@ async def _create_error(
             wrapper_elements=self._error_wrapper_elements(),
             status=response.status,
             context=context,
+            retry_after=parse_retry_after(response),
         )
 
     def _action_name(

packages/smithy-aws-core/src/smithy_aws_core/utils.py

@@ -1,7 +1,34 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
+import logging
+
 from smithy_core.documents import Document
 from smithy_core.shapes import ShapeID, ShapeType
+from smithy_http.aio.interfaces import HTTPResponse
+
+_LOGGER = logging.getLogger(__name__)
+
+_RETRY_AFTER_HEADER = "x-amz-retry-after"
+
+
+def parse_retry_after(response: HTTPResponse) -> float | None:
+    """Parse the ``x-amz-retry-after`` header into a backoff duration in seconds.
+
+    The header value is an integer number of milliseconds. Invalid or missing
+    values are ignored (return ``None``) so they fall back to exponential backoff.
+    """
+    if _RETRY_AFTER_HEADER not in response.fields:
+        return None
+    raw = response.fields[_RETRY_AFTER_HEADER].as_string()
+    try:
+        milliseconds = int(raw)
+    except (ValueError, TypeError):
+        _LOGGER.debug("Ignoring invalid %s header value: %r", _RETRY_AFTER_HEADER, raw)
+        return None
+    if milliseconds < 0:
+        _LOGGER.debug("Ignoring negative %s header value: %r", _RETRY_AFTER_HEADER, raw)
+        return None
+    return milliseconds / 1000.0
 
 
 def parse_document_discriminator(