Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ public final class AwsDynamoDbRetryIntegration implements PythonIntegration {


def dynamodb_retry_plugin(config: $1T) -> None:
\"\"\"Apply DynamoDB's standard-mode retry defaults for any option left unset.\"\"\"
\"\"\"Apply DynamoDB's standard and adaptive mode retry defaults for any option left unset.\"\"\"
retry_strategy = config.retry_strategy
if retry_strategy is not None and not isinstance(
retry_strategy, $2T
Expand All @@ -49,10 +49,11 @@ def dynamodb_retry_plugin(config: $1T) -> None:
if config.source_of("max_attempts") == $4T.DEFAULT:
max_attempts = None

if retry_mode != "standard":
if retry_mode not in ("standard", "adaptive"):
return

config.retry_strategy = $3T(
strategy_cls = $7T if retry_mode == "adaptive" else $3T
config.retry_strategy = strategy_cls(
max_attempts=(
max_attempts
if max_attempts is not None
Expand Down Expand Up @@ -86,6 +87,10 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
.namespace("smithy_core.aio.retries", ".")
.name("StandardRetryStrategy")
.build();
final Symbol adaptiveRetryStrategy = Symbol.builder()
.namespace("smithy_core.aio.retries", ".")
.name("AdaptiveRetryStrategy")
.build();
final Symbol configSource = Symbol.builder()
.namespace("smithy_aws_core.config", ".")
.name("ConfigSource")
Expand Down Expand Up @@ -122,7 +127,8 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
standardRetryStrategy,
configSource,
exponentialBackoffStrategy,
exponentialBackoffJitterType);
exponentialBackoffJitterType,
adaptiveRetryStrategy);
});
return List.of(filename);
})
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"type": "enhancement",
"description": "The `AWS_RETRY_MODE` environment variable and `retry_mode` config setting now accept `adaptive` instead of falling back to `standard` with a warning."
}
Original file line number Diff line number Diff line change
Expand Up @@ -78,10 +78,10 @@ class AsyncAwsConfig:
"""

retry_mode: RetryStrategyType | None = None
"""The retry mode to use. ``standard`` is the only accepted override.
"""The retry mode to use. Accepted overrides are ``standard`` and ``adaptive``.

``legacy`` and ``adaptive`` are rejected when set here; when they come from
the environment or a config file they warn and fall back to ``standard``.
``legacy`` is rejected when set here; when it comes from the environment or a
config file it warns and falls back to ``standard``.
"""

max_attempts: int | None = None
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -100,13 +100,6 @@ async def resolve_retry_mode(ctx: SharedConfigContext) -> Resolved[str | None]:
stacklevel=2,
)
return Resolved(value="standard", source=result.source)
elif result.value == "adaptive":
warnings.warn(
"'adaptive' retry mode is not supported, using 'standard' instead.",
UserWarning,
stacklevel=2,
)
return Resolved(value="standard", source=result.source)
return result


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,16 +36,15 @@ def validate_region(value: object) -> None:
def validate_retry_mode(retry_mode: str):
"""Validate retry mode.

Valid values: 'standard'
Valid values: 'standard', 'adaptive'

:param retry_mode: The retry mode value to validate
:raises: ConfigValidationError: If the retry mode is invalid
"""
# NOTE: RetryStrategyType includes 'simple' for direct config use, but the only valid
# string mode accepted here is 'standard'. 'adaptive' and 'legacy' are intentionally
# rejected as direct overrides ('adaptive' support may be added later; 'legacy' is not
# recommended and never will be). When 'adaptive' or 'legacy' come from the environment
# or a config file, resolve_retry_mode() warns and maps them to 'standard'.
# NOTE: RetryStrategyType includes 'simple' for direct config use, but it isn't a
# valid string mode here. 'legacy' is intentionally rejected as a direct override
# (it is not recommended and never will be). When 'legacy' comes from the
# environment or a config file, resolve_retry_mode() warns and maps it to 'standard'.

all_modes = list(get_args(RetryStrategyType))
if "simple" in all_modes:
Expand Down
21 changes: 9 additions & 12 deletions packages/smithy-aws-core/tests/unit/config/test_resolver.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
"""

import os
import warnings
from copy import deepcopy
from dataclasses import dataclass
from inspect import signature
Expand Down Expand Up @@ -552,7 +553,7 @@ async def test_resolves_from_profile_when_no_env(self):
assert result.source == ConfigSource.PROFILE

@pytest.mark.asyncio
async def test_adaptive_from_profile_warns_and_maps_to_standard(self):
async def test_resolves_adaptive_from_profile(self):
fs = FakeFileSystem(
{"/fake/config": "[profile default]\nretry_mode = adaptive\n"}
)
Expand All @@ -562,12 +563,10 @@ async def test_adaptive_from_profile_warns_and_maps_to_standard(self):
config_file_path="/fake/config",
credentials_file_path="/fake/credentials",
)
with pytest.warns(
UserWarning,
match="'adaptive' retry mode is not supported, using 'standard' instead.",
):
with warnings.catch_warnings():
warnings.simplefilter("error")
result = await resolve_retry_mode(ctx)
assert result.value == "standard"
assert result.value == "adaptive"
assert result.source == ConfigSource.PROFILE

@pytest.mark.asyncio
Expand Down Expand Up @@ -606,15 +605,13 @@ async def test_legacy_warns_and_maps_to_standard(self):
assert result.source == ConfigSource.ENV

@pytest.mark.asyncio
async def test_adaptive_warns_and_maps_to_standard(self):
async def test_resolves_adaptive_from_env(self):
with patch.dict(os.environ, {"AWS_RETRY_MODE": "adaptive"}, clear=True):
ctx = SharedConfigContext()
with pytest.warns(
UserWarning,
match="'adaptive' retry mode is not supported, using 'standard' instead.",
):
with warnings.catch_warnings():
warnings.simplefilter("error")
result = await resolve_retry_mode(ctx)
assert result.value == "standard"
assert result.value == "adaptive"
assert result.source == ConfigSource.ENV


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ def test_neither_default_nor_factory_raises(self):


class TestValidateRetryMode:
@pytest.mark.parametrize("mode", ["standard"])
@pytest.mark.parametrize("mode", ["standard", "adaptive"])
def test_valid_modes(self, mode: str):
validate_retry_mode(mode)

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"type": "bugfix",
"description": "`RetryStrategyResolver` now caches retry strategies per resolver instead of in a shared process-wide cache, so a client's retry state is no longer reset or shared when many clients are created."
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"type": "feature",
"description": "Added support for `adaptive` retry mode, which builds on standard retries by adding client-side rate limiting that slows the client's sending rate when a service throttles requests."
}
99 changes: 89 additions & 10 deletions packages/smithy-core/src/smithy_core/aio/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@
from .interfaces.auth import AuthScheme
from .interfaces.eventstream import EventReceiver
from .interfaces.retries import RetryStrategy
from .retries import RetryStrategyResolver
from .retries import AdaptiveRetryStrategy, RetryStrategyResolver
from .utils import close, seek

if TYPE_CHECKING:
Expand Down Expand Up @@ -384,27 +384,44 @@ async def _retry[I: SerializeableShape, O: DeserializeableShape](
request_future: Future[RequestContext[I, TRequest]] | None,
) -> OutputContext[I, O, TRequest | None, TResponse | None]:
if not call.retryable():
return await self._handle_attempt(call, request_context, request_future)
return await self._send_with_rate_limiting(
call, request_context, request_future
)

retry_strategy = call.retry_strategy
retry_token = await retry_strategy.acquire_initial_retry_token(
token_scope=call.retry_scope
)

last_output_context: OutputContext[I, O, TRequest, TResponse | None] | None = (
None
)
while True:
if retry_token.retry_delay:
await sleep(retry_token.retry_delay)

output_context = await self._handle_attempt(
call,
replace(
request_context,
transport_request=copy(request_context.transport_request),
),
request_future,
)
try:
output_context = await self._send_with_rate_limiting(
call,
replace(
request_context,
transport_request=copy(request_context.transport_request),
),
request_future,
)
except TimeoutError as e:
# Only the adaptive send-token wait raises this. If it happens on a
# retry, return the previous failed attempt (as when attempts run out)
# so the caller's error still includes that response's request ID.
if last_output_context is None:
raise
last_error = last_output_context.response
if isinstance(last_error, Exception):
last_error.add_note(f"Retry not attempted: {e}")
return last_output_context

if isinstance(output_context.response, Exception):
last_output_context = output_context
try:
retry_token = await retry_strategy.refresh_retry_token_for_retry(
token_to_renew=retry_token,
Expand Down Expand Up @@ -433,6 +450,68 @@ async def _retry[I: SerializeableShape, O: DeserializeableShape](
await retry_strategy.record_success(token=retry_token)
return output_context

async def _send_with_rate_limiting[
I: SerializeableShape,
O: DeserializeableShape,
](
self,
call: ClientCall[I, O],
request_context: RequestContext[I, TRequest],
request_future: Future[RequestContext[I, TRequest]] | None,
) -> OutputContext[I, O, TRequest, TResponse | None]:
"""Send a single attempt, applying client-side rate limiting (adaptive only).

Every attempt, retryable or not, acquires a send token before it is sent. Once a
response comes back from the service, the rate limiter is updated with it.
"""
retry_strategy = call.retry_strategy
try:
await self._handle_pre_request_rate_limiting(retry_strategy)
except Exception as e:
if request_future is not None and not request_future.done():
request_future.set_exception(e)
raise

output_context = await self._handle_attempt(
call, request_context, request_future
)

if output_context.transport_response is None:
# Failed before sending, or the connection failed: not a service response
return output_context

if isinstance(output_context.response, Exception):
await self._handle_post_error_response_rate_limiting(
retry_strategy, output_context.response
)
else:
await self._handle_success_rate_limiting(retry_strategy)
return output_context

async def _handle_pre_request_rate_limiting(
self, retry_strategy: RetryStrategy
) -> None:
"""Handle rate limiting before sending request."""
if isinstance(retry_strategy, AdaptiveRetryStrategy):
await retry_strategy.acquire_from_token_bucket()

async def _handle_post_error_response_rate_limiting(
self, retry_strategy: RetryStrategy, error: Exception
) -> None:
"""Handle rate limiting after failed response."""
if isinstance(retry_strategy, AdaptiveRetryStrategy):
is_throttling = retry_strategy.is_throttling_error(error)
await retry_strategy.rate_limiter.after_receiving_response(is_throttling)

async def _handle_success_rate_limiting(
self, retry_strategy: RetryStrategy
) -> None:
"""Handle rate limiting after successful response."""
if isinstance(retry_strategy, AdaptiveRetryStrategy):
await retry_strategy.rate_limiter.after_receiving_response(
throttling_error=False
)

async def _handle_attempt[I: SerializeableShape, O: DeserializeableShape](
self,
call: ClientCall[I, O],
Expand Down
Loading
Loading