Implement updated standard retry behavior

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

@@ -0,0 +1,42 @@
+/*
+ * 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.codegen.core.Symbol;
+import software.amazon.smithy.codegen.core.SymbolReference;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
+import software.amazon.smithy.python.codegen.integrations.RuntimeClientPlugin;
+
+/**
+ * Registers DynamoDB's retry defaults as a service-scoped client plugin on
+ * DynamoDB / DynamoDB Streams clients.
+ */
+public final class AwsDynamoDbRetryIntegration implements PythonIntegration {
+
+    private static final Set<String> DYNAMODB_SDK_IDS = Set.of("DynamoDB", "DynamoDB Streams");
+
+    private static final SymbolReference DYNAMODB_RETRY_PLUGIN = SymbolReference.builder()
+            .symbol(Symbol.builder()
+                    .namespace("smithy_aws_core.plugins", ".")
+                    .name("dynamodb_retry_plugin")
+                    .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
+                    .build())
+            .build();
+
+    @Override
+    public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
+        return List.of(
+                RuntimeClientPlugin.builder()
+                        .servicePredicate((model, service) -> service.getTrait(ServiceTrait.class)
+                                .map(trait -> DYNAMODB_SDK_IDS.contains(trait.getSdkId()))
+                                .orElse(false))
+                        .pythonPlugin(DYNAMODB_RETRY_PLUGIN)
+                        .build());
+    }
+}

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/.changes/next-release/smithy-aws-core-feature-9b04d7ee0977489ba85b40e929732ba0.json

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "Added DynamoDB and DynamoDB Streams retry defaults (4 max attempts and a 25ms non-throttling backoff)."
+}

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/plugins/__init__.py

@@ -0,0 +1,36 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+"""Service-level client plugins for AWS clients."""
+
+from typing import Protocol
+
+from smithy_core.aio.interfaces.retries import RetryStrategy
+from smithy_core.aio.retries import StandardRetryStrategy
+from smithy_core.retries import RetryStrategyOptions
+
+# DynamoDB uses max_attempts 4 and 25ms non-throttling backoff.
+_DYNAMODB_DEFAULT_MAX_ATTEMPTS = 4
+_DYNAMODB_DEFAULT_BACKOFF_SCALE = 0.025
+
+
+class _RetryConfig(Protocol):
+    retry_strategy: RetryStrategy | RetryStrategyOptions | None
+
+
+def dynamodb_retry_plugin(config: _RetryConfig) -> None:
+    """Apply DynamoDB's standard-mode retry defaults for any option left unset."""
+    rs = config.retry_strategy
+    if rs is not None and not isinstance(rs, RetryStrategyOptions):
+        return
+    if rs is not None and rs.retry_mode != "standard":
+        return
+
+    max_attempts = (
+        rs.max_attempts
+        if rs is not None and rs.max_attempts is not None
+        else _DYNAMODB_DEFAULT_MAX_ATTEMPTS
+    )
+    config.retry_strategy = StandardRetryStrategy(
+        max_attempts=max_attempts,
+        default_backoff_scale=_DYNAMODB_DEFAULT_BACKOFF_SCALE,
+    )