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,39 @@
+/*
+ * 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.python.codegen.integrations.PythonIntegration;
+
+/**
+ * Marks the known long-polling operations so the generic client generator
+ * applies long-polling retry behavior to them.
+ *
+ * <p>These operations are hard-coded until the {@code aws.api#longPoll} trait
+ * ships in service models. Once it ships, this can check for the trait instead.
+ */
+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) {
+        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("""
@@ -240,7 +258,8 @@ private void writeSharedOperationInit(
                     auth_scheme_resolver=config.auth_scheme_resolver,
                     supported_auth_schemes=config.auth_schemes,
                     endpoint_resolver=config.endpoint_resolver,
-                    retry_strategy=retry_strategy,
+                    retry_strategy=retry_strategy,${?isLongPolling}
+                    is_long_polling=True,${/isLongPolling}
                 )
                 """,
                 writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)),

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,12 +32,15 @@ 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, is_long_polling: bool = False
     ) -> 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 is_long_polling: Whether the operation is a long-polling operation.
+            Long-polling operations must back off before returning even when the
+            retry quota is exhausted.
         :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.
@@ -110,8 +113,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 in the `awsQuery` protocol."
+}

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,7 +27,11 @@
 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,
+    parse_retry_after,
+)
 from smithy_http.deserializers import HTTPResponseDeserializer
 
 from .._private.query.errors import (
@@ -353,6 +357,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/tests/unit/test_query.py

@@ -2,13 +2,18 @@
 #  SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass
 from io import BytesIO
+from typing import Any, cast
+from unittest.mock import Mock
 
+from smithy_aws_core._private.query.errors import create_aws_query_error
 from smithy_aws_core._private.query.serializers import QueryShapeSerializer
+from smithy_core.documents import TypeRegistry
 from smithy_core.prelude import STRING
-from smithy_core.schemas import Schema
+from smithy_core.schemas import APIOperation, Schema
 from smithy_core.serializers import ShapeSerializer
 from smithy_core.shapes import ShapeID, ShapeType
 from smithy_core.traits import XMLFlattenedTrait, XMLNameTrait
+from smithy_core.types import TypedProperties
 
 
 def test_query_list_serialization() -> None:
@@ -280,3 +285,39 @@ def serialize_members(self, serializer: ShapeSerializer) -> None:
     Outer(inner=Inner("x")).serialize(serializer)
 
     assert params == [("inner.value", "x")]
+
+
+def _error_test_operation() -> APIOperation[Any, Any]:
+    operation = Mock(spec=APIOperation)
+    operation.schema = Schema(
+        id=ShapeID("com.example#TestOp"), shape_type=ShapeType.OPERATION
+    )
+    operation.error_schemas = []
+    return cast("APIOperation[Any, Any]", operation)
+
+
+def test_aws_query_error_sets_retry_after_on_generic_error() -> None:
+    error = create_aws_query_error(
+        body=b"",
+        operation=_error_test_operation(),
+        error_registry=TypeRegistry({}),
+        default_namespace="com.example",
+        wrapper_elements=("ErrorResponse", "Error"),
+        status=503,
+        context=TypedProperties(),
+        retry_after=1.5,
+    )
+    assert error.retry_after == 1.5
+
+
+def test_aws_query_error_retry_after_none_by_default() -> None:
+    error = create_aws_query_error(
+        body=b"",
+        operation=_error_test_operation(),
+        error_registry=TypeRegistry({}),
+        default_namespace="com.example",
+        wrapper_elements=("ErrorResponse", "Error"),
+        status=503,
+        context=TypedProperties(),
+    )
+    assert error.retry_after is None

packages/smithy-core/.changes/next-release/smithy-core-breaking-dc093af32db648139ed3f64cf4251a40.json

@@ -0,0 +1,4 @@
+{
+  "type": "breaking",
+  "description": "Updated retry quota costs and added an `is_long_polling` argument to `acquire_initial_retry_token`."
+}

packages/smithy-core/.changes/next-release/smithy-core-feature-aa633008af584d94ad03679f73a4d436.json

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "Added support for separate backoff for throttling errors, the `x-amz-retry-after` header, per-service retry defaults, and long-polling backoff."
+}