diff --git a/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/customizations/dynamodb/AwsDynamoDbRetryIntegration.java b/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/customizations/dynamodb/AwsDynamoDbRetryIntegration.java index 39d74f324..6b225a4c0 100644 --- a/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/customizations/dynamodb/AwsDynamoDbRetryIntegration.java +++ b/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/customizations/dynamodb/AwsDynamoDbRetryIntegration.java @@ -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 @@ -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 @@ -86,6 +87,10 @@ public List 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") @@ -122,7 +127,8 @@ public List getClientPlugins(GenerationContext context) { standardRetryStrategy, configSource, exponentialBackoffStrategy, - exponentialBackoffJitterType); + exponentialBackoffJitterType, + adaptiveRetryStrategy); }); return List.of(filename); }) diff --git a/packages/smithy-aws-core/.changes/next-release/smithy-aws-core-enhancement-283eae54d51246a6b6399e6c0963fdb9.json b/packages/smithy-aws-core/.changes/next-release/smithy-aws-core-enhancement-283eae54d51246a6b6399e6c0963fdb9.json new file mode 100644 index 000000000..3d2e858b0 --- /dev/null +++ b/packages/smithy-aws-core/.changes/next-release/smithy-aws-core-enhancement-283eae54d51246a6b6399e6c0963fdb9.json @@ -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." +} \ No newline at end of file diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py b/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py index bbd80eea0..f6ee2c7bd 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py @@ -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 diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py b/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py index a2f03c4d1..31b363876 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py @@ -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 diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/validators.py b/packages/smithy-aws-core/src/smithy_aws_core/config/validators.py index ea4d018f0..376736a26 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/validators.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/validators.py @@ -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: diff --git a/packages/smithy-aws-core/tests/unit/config/test_resolver.py b/packages/smithy-aws-core/tests/unit/config/test_resolver.py index 7d09ade96..04014d94b 100644 --- a/packages/smithy-aws-core/tests/unit/config/test_resolver.py +++ b/packages/smithy-aws-core/tests/unit/config/test_resolver.py @@ -7,6 +7,7 @@ """ import os +import warnings from copy import deepcopy from dataclasses import dataclass from inspect import signature @@ -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"} ) @@ -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 @@ -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 diff --git a/packages/smithy-aws-core/tests/unit/config/test_validators.py b/packages/smithy-aws-core/tests/unit/config/test_validators.py index 2b77825c8..4c1015943 100644 --- a/packages/smithy-aws-core/tests/unit/config/test_validators.py +++ b/packages/smithy-aws-core/tests/unit/config/test_validators.py @@ -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) diff --git a/packages/smithy-core/.changes/next-release/smithy-core-bugfix-13759af1cbc74199b4cce812e99b77e8.json b/packages/smithy-core/.changes/next-release/smithy-core-bugfix-13759af1cbc74199b4cce812e99b77e8.json new file mode 100644 index 000000000..1846fadf7 --- /dev/null +++ b/packages/smithy-core/.changes/next-release/smithy-core-bugfix-13759af1cbc74199b4cce812e99b77e8.json @@ -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." +} \ No newline at end of file diff --git a/packages/smithy-core/.changes/next-release/smithy-core-feature-945068e3c48d45038871c73459cf16c5.json b/packages/smithy-core/.changes/next-release/smithy-core-feature-945068e3c48d45038871c73459cf16c5.json new file mode 100644 index 000000000..f6049f97d --- /dev/null +++ b/packages/smithy-core/.changes/next-release/smithy-core-feature-945068e3c48d45038871c73459cf16c5.json @@ -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." +} \ No newline at end of file diff --git a/packages/smithy-core/src/smithy_core/aio/client.py b/packages/smithy-core/src/smithy_core/aio/client.py index c358a9c1b..7afadb608 100644 --- a/packages/smithy-core/src/smithy_core/aio/client.py +++ b/packages/smithy-core/src/smithy_core/aio/client.py @@ -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: @@ -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, @@ -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], diff --git a/packages/smithy-core/src/smithy_core/aio/retries.py b/packages/smithy-core/src/smithy_core/aio/retries.py index 46c172b2c..47faf3a2b 100644 --- a/packages/smithy-core/src/smithy_core/aio/retries.py +++ b/packages/smithy-core/src/smithy_core/aio/retries.py @@ -1,11 +1,16 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from functools import lru_cache +import asyncio +import math +import threading +import time +import weakref from typing import Any from ..exceptions import RetryError from ..interfaces import retries as retries_interface from ..retries import ( + CubicCalculator, ExponentialBackoffJitterType, ExponentialRetryBackoffStrategy, RetryStrategyOptions, @@ -21,9 +26,14 @@ class RetryStrategyResolver: """Retry strategy resolver that caches retry strategies based on configuration options. This resolver caches retry strategy instances based on their configuration to reuse existing - instances of RetryStrategy with the same settings. Uses LRU cache for thread-safe caching. + instances of RetryStrategy with the same settings. Each resolver has its own cache, so a + client keeps its strategies, and any state they hold such as an adaptive rate limiter, + for as long as the client exists, no matter how many other clients are created. """ + def __init__(self) -> None: + self._strategies: dict[tuple[RetryStrategyType, int | None], RetryStrategy] = {} + async def resolve_retry_strategy( self, *, @@ -57,9 +67,16 @@ async def resolve_retry_strategy( retry_strategy.retry_mode, retry_strategy.max_attempts ) - @lru_cache def _create_retry_strategy( self, retry_mode: RetryStrategyType, max_attempts: int | None + ) -> RetryStrategy: + key = (retry_mode, max_attempts) + if key not in self._strategies: + self._strategies[key] = self._new_retry_strategy(retry_mode, max_attempts) + return self._strategies[key] + + def _new_retry_strategy( + self, retry_mode: RetryStrategyType, max_attempts: int | None ) -> RetryStrategy: kwargs: dict[str, Any] = {"max_attempts": max_attempts} filtered_kwargs: dict[str, Any] = { @@ -70,6 +87,8 @@ def _create_retry_strategy( return SimpleRetryStrategy(**filtered_kwargs) case "standard": return StandardRetryStrategy(**filtered_kwargs) + case "adaptive": + return AdaptiveRetryStrategy(**filtered_kwargs) case _: raise ValueError(f"Unknown retry mode: {retry_mode}") @@ -272,3 +291,400 @@ async def record_success(self, *, token: retries_interface.RetryToken) -> None: def __deepcopy__(self, memo: Any) -> "StandardRetryStrategy": return self + + +class TokenBucket: + """Token bucket for rate limiting with configurable fill rate. + + TokenBucket provides a collection of arbitrary tokens while managing issuance + and refilling over time. This is controlled by a fill rate that can be variably + adjusted. When tokens aren't available, the bucket will enforce a delay before + attempting to reacquire tokens until one is available or the defined timeout is + reached. + + The bucket is safe to share across threads, each running its own event loop. A + threading.Lock guards the token count and rate, because an asyncio primitive only + coordinates tasks on a single event loop. The lock only covers short arithmetic and + is never held across an ``await``, so it never blocks for long. Waiting for tokens + uses a per-event-loop asyncio.Condition and does not hold the lock. + """ + + MIN_FILL_RATE = 0.5 # Minimum allowed fill rate (0.5 tokens/second) + MIN_CAPACITY = 1.0 # Minimum allowed bucket capacity (1.0 tokens) + DEFAULT_TIMEOUT = 30.0 # Default timeout for token acquisition (30.0 seconds) + + def __init__( + self, + *, + curr_capacity: float = 0, + fill_rate: float = MIN_FILL_RATE, + timeout: float = DEFAULT_TIMEOUT, + ): + """Initialize a new TokenBucket. + + :param curr_capacity: Initial number of tokens in the bucket. The bucket can + hold at least MIN_CAPACITY tokens, but starts empty by default. + :param fill_rate: Rate at which tokens are added to the bucket (tokens/second). + :param timeout: Maximum time to wait for token acquisition before + raising TimeoutError. + """ + self._curr_capacity: float = curr_capacity + self._max_capacity: float = max(curr_capacity, self.MIN_CAPACITY) + self._fill_rate: float = max(fill_rate, self.MIN_FILL_RATE) + self._timeout = timeout + self._last_timestamp: float = time.monotonic() + self._lock = threading.Lock() + self._conditions: weakref.WeakKeyDictionary[ + asyncio.AbstractEventLoop, asyncio.Condition + ] = weakref.WeakKeyDictionary() + + async def acquire(self, amount: float) -> None: + """Acquire tokens from the bucket. + + If sufficient tokens are available, they are immediately consumed and the + method returns. If insufficient tokens are available, the method will wait + until enough tokens have been refilled or the timeout is reached. + + :param amount: Number of tokens to acquire. + + :raises TimeoutError: Acquisition took longer than the configured timeout. + """ + start_time = time.monotonic() + condition = self._get_condition() + async with condition: + while True: + with self._lock: + self._refill() + if self._curr_capacity >= amount: + self._curr_capacity -= amount + return + wait_time = (amount - self._curr_capacity) / self._fill_rate + + elapsed = time.monotonic() - start_time + if elapsed >= self._timeout: + # Fails the request; this is not counted as a retry attempt + raise TimeoutError( + f"Failed to acquire {amount} tokens within {self._timeout}s" + ) + # Wait outside the thread lock, and never past the timeout + try: + await asyncio.wait_for( + condition.wait(), + timeout=min(wait_time, self._timeout - elapsed), + ) + except TimeoutError: + pass + + def _get_condition(self) -> asyncio.Condition: + """Return the Condition for the event loop that is currently running. + + An asyncio.Condition belongs to the first event loop that waits on it, and + using it from any other loop raises a RuntimeError. The bucket lives as long + as the client, so it can outlive the loop it was first used on. + + For example, a client using the CRT HTTP client can be reused across + separate asyncio.run() calls, such as one call per Lambda invocation. If a + request waits for a token during the first call, the Condition is tied to + that loop. Without this check, every request that has to wait for a token + in a later asyncio.run() call would fail. + + To avoid that, each event loop gets its own Condition, which is reused for as + long as that loop exists and is released along with it. + """ + loop = asyncio.get_running_loop() + with self._lock: + condition = self._conditions.get(loop) + if condition is None: + condition = asyncio.Condition() + self._conditions[loop] = condition + return condition + + def _refill(self) -> None: + curr_time = time.monotonic() + elapsed = curr_time - self._last_timestamp + refill_amount = elapsed * self._fill_rate + self._curr_capacity = min( + self._max_capacity, self._curr_capacity + refill_amount + ) + self._last_timestamp = curr_time + + async def update_rate(self, rate: float) -> None: + """Update the bucket's fill rate, maximum capacity and current capacity (if its + greater than maximum capacity). + + :param rate: New fill rate (tokens/second). It won't be less than MIN_FILL_RATE. + Current capacity will be reduced if it exceeds the new maximum capacity. + """ + self._set_rate(rate) + await self._notify_waiters() + + def _set_rate(self, rate: float) -> None: + with self._lock: + self._refill() + self._fill_rate = max(rate, self.MIN_FILL_RATE) + self._max_capacity = max(rate, self.MIN_CAPACITY) + self._curr_capacity = min(self._curr_capacity, self._max_capacity) + + async def _notify_waiters(self) -> None: + # Wakes waiters on this event loop. Waiters on other threads' loops pick up + # the new rate when their current wait times out. + condition = self._get_condition() + async with condition: + condition.notify_all() + + @property + def current_capacity(self) -> float: + """Get the current number of tokens in the bucket. + + :return: The current token count as of the last refill operation. + """ + return self._curr_capacity + + @property + def max_capacity(self) -> float: + """Get the maximum capacity of the bucket. + + :return: The maximum number of tokens the bucket can hold. + """ + return self._max_capacity + + @property + def fill_rate(self) -> float: + """Get the current fill rate of the bucket. + + :return: The rate at which tokens are added to the bucket (tokens/second). + """ + return self._fill_rate + + +class RequestRateTracker: + """RequestRateTracker tracks the client's request sending rate. + + RequestRateTracker measures the actual client request sending rate using + time-bucketed sampling with exponential smoothing. It tracks requests in + half-second intervals by default and calculates a smoothed average rate + to provide accurate measurements for adaptive rate limiting algorithms. + """ + + _DEFAULT_SMOOTHING = 0.8 + _TIME_BUCKET_RANGE = 0.5 + + def __init__( + self, + smoothing: float = _DEFAULT_SMOOTHING, + time_bucket_range: float = _TIME_BUCKET_RANGE, + ): + """Initialize a new RequestRateTracker. + + :param smoothing: Exponential smoothing factor. This constant + represents how much weight is given to recent measurements. + Higher values place more emphasis on the most recent observations. + :param time_bucket_range: Time bucket duration in seconds. Rate + calculations are updated when transitioning between buckets. + """ + self._smoothing = smoothing + self._time_bucket_scale = 1 / time_bucket_range + self._request_count = 0 + self._last_calculated_time_bucket = math.floor(time.monotonic()) + self._measured_rate = 0 + # The update never awaits, so it only needs protecting across threads + self._lock = threading.Lock() + + async def measure_rate(self) -> float: + """Measure and return the current request rate. + + Increments the request count and calculates a new smoothed rate when + transitioning to a new time bucket. Returns the current measured rate + without recalculation if still within the same time bucket. + :return: Current smoothed request rate in requests per second. + """ + with self._lock: + curr_time = time.monotonic() + current_time_bucket = ( + math.floor(curr_time * self._time_bucket_scale) + / self._time_bucket_scale + ) + self._request_count += 1 + if current_time_bucket > self._last_calculated_time_bucket: + current_rate = self._request_count / ( + current_time_bucket - self._last_calculated_time_bucket + ) + self._measured_rate = (current_rate * self._smoothing) + ( + self._measured_rate * (1 - self._smoothing) + ) + self._request_count = 0 + self._last_calculated_time_bucket = current_time_bucket + return self._measured_rate + + @property + def request_count(self) -> int: + """Get current request count. For testing only.""" + return self._request_count + + @property + def measured_client_rate(self) -> float: + """Get the client's sending rate. For testing only.""" + return self._measured_rate + + +class ClientRateLimiter: + """ClientRateLimiter limits the rate of requests. + + ClientRateLimiter implements adaptive rate limiting using token bucket and CUBIC + algorithm. It controls request sending rates by acquiring tokens before requests + and dynamically adjusting rates based on service responses - reducing rates when + throttled and increasing rates during successful periods to optimize throughput + while preventing service overload. + """ + + _REQUEST_COST = 1.0 + + def __init__( + self, + token_bucket: TokenBucket, + cubic_calculator: CubicCalculator, + rate_tracker: RequestRateTracker, + rate_limiter_enabled: bool, + ): + """Initialize a new ClientRateLimiter. + + :param token_bucket: Token bucket for controlling request sending rates. + :param cubic_calculator: CUBIC algorithm calculator for rate adjustments. + :param rate_tracker: Tracker for measuring actual client request rates. + :param rate_limiter_enabled: Whether rate limiting is enabled. + """ + + self._rate_tracker = rate_tracker + self._token_bucket = token_bucket + self._cubic_calculator = cubic_calculator + self._rate_limiter_enabled = rate_limiter_enabled + self._lock = threading.Lock() + + async def before_sending_request(self) -> None: + """Acquire a token before making a request.""" + if self._rate_limiter_enabled: + await self._token_bucket.acquire(self._REQUEST_COST) + + async def after_receiving_response(self, throttling_error: bool) -> None: + """Update the request rate based on the response using CUBIC algorithm. + + Reduces the rate by 30% when throttled, or increases the rate using + CUBIC scaling for successful responses. Updates the token bucket with + the new calculated rate, capped at 2x the measured client rate. + :param throttling_error: True if the response was a throttling error. + """ + measured_rate = await self._rate_tracker.measure_rate() + with self._lock: + timestamp = time.monotonic() + if throttling_error: + if not self._rate_limiter_enabled: + rate_to_use = measured_rate + else: + fill_rate = self._token_bucket.fill_rate + rate_to_use = min(measured_rate, fill_rate) + + cubic_calculated_rate = ( + self._cubic_calculator.calculate_throttled_request_rate( + rate_to_use, timestamp + ) + ) + self._rate_limiter_enabled = True + else: + cubic_calculated_rate = ( + self._cubic_calculator.calculate_scaled_request_rate(timestamp) + ) + + new_rate = min(cubic_calculated_rate, 2 * measured_rate) + # Set under the lock so rates reach the bucket in the order they're + # calculated. Waking waiters awaits, so it happens after the lock. + self._token_bucket._set_rate(new_rate) # pyright: ignore[reportPrivateUsage] + await self._token_bucket._notify_waiters() # pyright: ignore[reportPrivateUsage] + + @property + def rate_limit_enabled(self) -> bool: + return self._rate_limiter_enabled + + +class AdaptiveRetryStrategy(StandardRetryStrategy): + """Adaptive retry strategy with client-side rate limiting using CUBIC algorithm. + + Builds on top of StandardRetryStrategy by adding token bucket rate limiting and + CUBIC congestion control. Rate limiting is enabled after the first throttling + response and dynamically adjusts sending rates based on the response type. + + Use this strategy when a client sends a high volume of requests to a service and + frequently receives throttling errors. Rather than only backing off individual + retries, it slows the client's overall sending rate, reducing requests that are + throttled and wasted on retries. + """ + + STARTING_MAX_RATE = 0 + + def __init__( + self, + *, + backoff_strategy: retries_interface.RetryBackoffStrategy | None = None, + throttling_backoff_strategy: retries_interface.RetryBackoffStrategy + | None = None, + max_attempts: int = 3, + retry_quota: StandardRetryQuota | None = None, + rate_limiter: ClientRateLimiter | None = None, + ): + """Initialize AdaptiveRetryStrategy. + + :param backoff_strategy: The backoff strategy used to compute the retry delay + for non-throttling errors. Defaults to a 50ms-base + :py:class:`ExponentialRetryBackoffStrategy`. + + :param throttling_backoff_strategy: The backoff strategy used to compute the + retry delay for throttling errors. Defaults to a 1000ms-base + :py:class:`ExponentialRetryBackoffStrategy`. + + :param max_attempts: Upper limit on total number of attempts made, including + initial attempt and retries. + + :param retry_quota: The retry quota to use for managing retry capacity. Defaults + to a new :py:class:`StandardRetryQuota` instance. + + :param rate_limiter: Optional pre-configured rate limiter. If None, creates + default components with rate limiting initially disabled. + """ + super().__init__( + backoff_strategy=backoff_strategy, + throttling_backoff_strategy=throttling_backoff_strategy, + max_attempts=max_attempts, + retry_quota=retry_quota, + ) + + if rate_limiter is None: + token_bucket = TokenBucket() + cubic_calculator = CubicCalculator( + starting_max_rate=self.STARTING_MAX_RATE, start_time=time.monotonic() + ) + rate_tracker = RequestRateTracker() + self._rate_limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=cubic_calculator, + rate_tracker=rate_tracker, + rate_limiter_enabled=False, # Disabled until first throttle + ) + else: + self._rate_limiter = rate_limiter + + @property + def rate_limiter(self) -> ClientRateLimiter: + """Get the rate limiter for integration with request pipeline.""" + return self._rate_limiter + + def is_throttling_error(self, error: Exception) -> bool: + """Check if error is a throttling error using existing ErrorRetryInfo.""" + if isinstance(error, retries_interface.ErrorRetryInfo): + return error.is_throttling_error + return False + + async def acquire_from_token_bucket(self) -> None: + if self._rate_limiter.rate_limit_enabled: + await self._rate_limiter.before_sending_request() + + def __deepcopy__(self, memo: Any) -> "AdaptiveRetryStrategy": + return self diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index ed9bcf46a..85e1af60f 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import random import threading +import time from collections.abc import Callable from dataclasses import dataclass from enum import Enum @@ -10,7 +11,7 @@ from .exceptions import RetryError from .interfaces import retries as retries_interface -RetryStrategyType = Literal["simple", "standard"] +RetryStrategyType = Literal["simple", "standard", "adaptive"] @dataclass(kw_only=True, frozen=True) @@ -270,3 +271,88 @@ class StandardRetryToken: quota_acquired: int = 0 """The amount of quota acquired for this retry attempt.""" + + +class CubicCalculator: + """CubicCalculator calculates a new rate using a modified CUBIC algorithm. + + CubicCalculator implements the CUBIC congestion control algorithm for + adaptive rate limiting. It dynamically adjusts request rates based on + throttling responses, reducing rates by 30% when throttled and + gradually increasing rates using a scale function when the request + is successful. + """ + + # Scale constant used to scale up requests + _SCALE_CONSTANT = 0.4 + # Beta constant used to slow down requests + _BETA = 0.7 + + def __init__( + self, + starting_max_rate: float = 0, + start_time: float | None = None, + ): + """Initialize a new CubicCalculator. + + :param starting_max_rate: Initial maximum request per second. + :param start_time: Initial time of the CubicCalculator. + """ + if starting_max_rate < 0: + raise ValueError( + f"starting_max_rate must not be negative, got {starting_max_rate}" + ) + if start_time is None: + start_time = time.monotonic() + + self._last_max_rate = starting_max_rate + self._last_throttle_time = start_time + self._inflection_point_time = self.calculate_and_update_inflection_point() + + def calculate_and_update_inflection_point(self) -> float: + """Calculate and update the CUBIC inflection point for rate recovery. + + After throttling, the CUBIC curve returns to the previous maximum rate + after exactly `_inflection_point_time` seconds. Before this point, the + rate grows slowly. After this point, it grows rapidly. + :return: Inflection point time in seconds for the CUBIC algorithm. + """ + self._inflection_point_time = ( + (self._last_max_rate * (1 - self._BETA)) / self._SCALE_CONSTANT + ) ** (1 / 3.0) + return self._inflection_point_time + + def calculate_scaled_request_rate(self, timestamp: float) -> float: + """Scale up the request rate after a successful response. + + :param timestamp: Timestamp of the response. + :return: New calculated request rate based on CUBIC scaling. + """ + dt = timestamp - self._last_throttle_time + calculated_rate = ( + self._SCALE_CONSTANT * ((dt - self._inflection_point_time) ** 3) + ) + self._last_max_rate + return calculated_rate + + def calculate_throttled_request_rate( + self, rate_to_use: float, timestamp: float + ) -> float: + """Throttle the request rate after a throttled response is received. + + :param rate_to_use: Current request rate in use. + :param timestamp: Timestamp of the response. + :return: New calculated request rate based on CUBIC throttling. + """ + self._last_max_rate = rate_to_use + self.calculate_and_update_inflection_point() + calculated_rate = rate_to_use * self._BETA + self._last_throttle_time = timestamp + return calculated_rate + + @property + def last_max_rate(self) -> float: + return self._last_max_rate + + @property + def last_throttle_time(self) -> float: + return self._last_throttle_time diff --git a/packages/smithy-core/tests/functional/test_retries.py b/packages/smithy-core/tests/functional/test_retries.py index 8fbc3d1eb..04e7de6ff 100644 --- a/packages/smithy-core/tests/functional/test_retries.py +++ b/packages/smithy-core/tests/functional/test_retries.py @@ -1,17 +1,24 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from asyncio import gather, sleep +from asyncio import gather, sleep, wait_for +from collections.abc import Coroutine +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from smithy_core.aio.client import RequestPipeline from smithy_core.aio.interfaces import retries as retries_interface -from smithy_core.aio.retries import StandardRetryStrategy +from smithy_core.aio.retries import AdaptiveRetryStrategy, StandardRetryStrategy from smithy_core.exceptions import CallError, ClientTimeoutError, RetryError +from smithy_core.interceptors import RequestContext from smithy_core.retries import ( ExponentialBackoffJitterType, ExponentialRetryBackoffStrategy, StandardRetryQuota, ) +from smithy_core.types import TypedProperties # TODO: Refactor this to use a smithy-testing generated client @@ -151,3 +158,264 @@ async def test_retry_quota_handles_timeout_errors(): assert result == "success" assert attempts == 3 assert quota.available_capacity == 486 + + +class FakeClock: + """Simulated time for the adaptive retry pipeline. + + Replaces the clock read by the rate limiter, the retry backoff sleep, and the + token bucket wait, so simulated seconds pass instantly and deterministically. + """ + + # A real event loop always advances time by at least this much per wait + _MIN_WAIT = 1e-6 + + def __init__(self) -> None: + self.now = 0.0 + self.backoffs: list[float] = [] + + def monotonic(self) -> float: + return self.now + + async def sleep(self, delay: float) -> None: + self.backoffs.append(delay) + self.now += delay + + async def wait_for(self, coro: Coroutine[Any, Any, Any], timeout: float) -> None: # noqa: ASYNC109 + # Nothing notifies a waiter in a single-client simulation, so every wait + # times out after the full duration + coro.close() + self.now += max(timeout, self._MIN_WAIT) + raise TimeoutError + + +class TestAdaptiveRetryPipeline: + @pytest.fixture + def clock(self): + clock = FakeClock() + with ( + patch( + "smithy_core.aio.retries.time", + SimpleNamespace(monotonic=clock.monotonic), + ), + patch("smithy_core.aio.client.sleep", clock.sleep), + patch("smithy_core.aio.retries.asyncio.wait_for", clock.wait_for), + ): + yield clock + + def _call(self, strategy: AdaptiveRetryStrategy) -> Any: + call = MagicMock() + call.retryable.return_value = True + call.retry_strategy = strategy + call.retry_scope = None + return call + + def _request_context(self) -> RequestContext[Any, Any]: + return RequestContext( + request=MagicMock(), + transport_request=SimpleNamespace(body=b""), + properties=TypedProperties(), + ) + + def _response(self, status: int) -> CallError | str: + if status == 200: + return "ok" + return CallError( + fault="server" if status >= 500 else "client", + message=f"HTTP {status}", + is_retry_safe=True, + is_throttling_error=status == 429, + ) + + @pytest.mark.parametrize( + "statuses, expected", + [ + # (rate limiting wait, response status, backoff, terminal) + ( + [429, 429, 429], + [ + (0.0, 429, 0.5, False), + (1.5, 429, 1.0, False), + (1.0, 429, None, True), + ], + ), + ( + [429, 500, 200], + [ + (0.0, 429, 0.5, False), + (1.5, 500, 1.0, False), + (0.0, 200, None, True), + ], + ), + ], + ids=["throttled_until_max_attempts", "throttle_then_success"], + ) + async def test_adaptive_retry_sequence( + self, + clock: FakeClock, + statuses: list[int], + expected: list[tuple[float, int, float | None, bool]], + ): + # Full jitter that always picks the midpoint, so backoff is deterministic + backoff = ExponentialRetryBackoffStrategy( + backoff_scale_value=1, + jitter_type=ExponentialBackoffJitterType.FULL, + random=lambda: 0.5, + ) + strategy = AdaptiveRetryStrategy( + max_attempts=3, + backoff_strategy=backoff, + throttling_backoff_strategy=backoff, + ) + pipeline: RequestPipeline[Any, Any] = RequestPipeline( + protocol=MagicMock(), transport=MagicMock() + ) + responses = iter(statuses) + sent: list[tuple[float, int]] = [] + attempt_start = 0.0 + + async def handle_attempt(*args: Any) -> SimpleNamespace: + status = next(responses) + sent.append((clock.now - attempt_start, status)) + return SimpleNamespace( + transport_response=MagicMock(), response=self._response(status) + ) + + original_pre_request = pipeline._handle_pre_request_rate_limiting # type: ignore[reportPrivateUsage] + + async def pre_request(retry_strategy: Any) -> None: + nonlocal attempt_start + attempt_start = clock.now + await original_pre_request(retry_strategy) + + with ( + patch.object( + pipeline, "_handle_attempt", AsyncMock(side_effect=handle_attempt) + ), + patch.object(pipeline, "_handle_pre_request_rate_limiting", pre_request), + ): + try: + await pipeline._retry( # type: ignore[reportPrivateUsage] + self._call(strategy), self._request_context(), None + ) + except CallError: + pass + + assert strategy.rate_limiter.rate_limit_enabled + backoffs: list[float | None] = [*clock.backoffs, None] + actual = [ + (round(wait, 6), status, backoffs[i], i == len(sent) - 1) + for i, (wait, status) in enumerate(sent) + ] + assert actual == expected + + async def test_adaptive_retry_send_rate_converges_to_server_limit( + self, clock: FakeClock + ): + # The server lets through 10 requests per second and throttles the rest, + # while the client attempts 100 requests per second + server_rps = 10.0 + client_interval = 0.01 + duration = 120.0 + + # No retries, so backoff can't slow the client; only rate limiting can + strategy = AdaptiveRetryStrategy(max_attempts=1) + pipeline: RequestPipeline[Any, Any] = RequestPipeline( + protocol=MagicMock(), transport=MagicMock() + ) + server_tokens = server_rps + server_last_refill = 0.0 + sent: list[float] = [] + succeeded: list[float] = [] + + async def handle_attempt(*args: Any) -> SimpleNamespace: + nonlocal server_tokens, server_last_refill + elapsed = clock.now - server_last_refill + server_tokens = min(server_rps, server_tokens + elapsed * server_rps) + server_last_refill = clock.now + sent.append(clock.now) + if server_tokens >= 1: + server_tokens -= 1 + succeeded.append(clock.now) + return SimpleNamespace( + transport_response=MagicMock(), response=self._response(200) + ) + return SimpleNamespace( + transport_response=MagicMock(), response=self._response(429) + ) + + with patch.object( + pipeline, "_handle_attempt", AsyncMock(side_effect=handle_attempt) + ): + while clock.now < duration: + try: + await pipeline._retry( # type: ignore[reportPrivateUsage] + self._call(strategy), self._request_context(), None + ) + except CallError: + pass + clock.now += client_interval + + # Measure the steady state over the last 30 simulated seconds + window_start = duration - 30 + send_rate = sum(t >= window_start for t in sent) / 30 + success_rate = sum(t >= window_start for t in succeeded) / 30 + assert send_rate < 15 + # Converging must not over-throttle the client well below the server limit + assert success_rate > 8 + + async def test_input_stream_raises_when_send_token_acquisition_fails(self): + strategy = AdaptiveRetryStrategy() + pipeline: RequestPipeline[Any, Any] = RequestPipeline( + protocol=MagicMock(), transport=MagicMock() + ) + call = self._call(strategy) + # Event stream inputs aren't retryable + call.retryable.return_value = False + error = TimeoutError("Failed to acquire 1.0 tokens within 30.0s") + + with patch.object( + strategy, "acquire_from_token_bucket", AsyncMock(side_effect=error) + ): + # Bound the wait so a regression fails instead of hanging the suite + with pytest.raises(TimeoutError, match="Failed to acquire"): + await wait_for(pipeline.input_stream(call, Any), timeout=1) + + def _pipeline_with_429_attempts( + self, strategy: AdaptiveRetryStrategy, acquire_results: list[Exception | None] + ) -> tuple[RequestPipeline[Any, Any], AsyncMock, Any]: + pipeline: RequestPipeline[Any, Any] = RequestPipeline( + protocol=MagicMock(), transport=MagicMock() + ) + attempt = AsyncMock( + return_value=SimpleNamespace( + transport_response=MagicMock(), response=self._response(429) + ) + ) + patches = ( + patch.object(pipeline, "_handle_attempt", attempt), + patch.object( + strategy, + "acquire_from_token_bucket", + AsyncMock(side_effect=acquire_results), + ), + patch("smithy_core.aio.client.sleep", AsyncMock()), + ) + return pipeline, attempt, patches + + async def test_send_token_timeout_on_first_attempt_raises(self): + quota = StandardRetryQuota(initial_capacity=500) + strategy = AdaptiveRetryStrategy(max_attempts=3, retry_quota=quota) + pipeline, attempt, (p1, p2, p3) = self._pipeline_with_429_attempts( + strategy, [TimeoutError("Failed to acquire")] + ) + + with p1, p2, p3: + with pytest.raises(TimeoutError, match="Failed to acquire"): + await pipeline._retry( # type: ignore[reportPrivateUsage] + self._call(strategy), self._request_context(), None + ) + + # Nothing was sent, so no retry quota was spent + attempt.assert_not_awaited() + assert quota.available_capacity == 500 diff --git a/packages/smithy-core/tests/unit/aio/test_client.py b/packages/smithy-core/tests/unit/aio/test_client.py index f971d142b..81c69b896 100644 --- a/packages/smithy-core/tests/unit/aio/test_client.py +++ b/packages/smithy-core/tests/unit/aio/test_client.py @@ -3,12 +3,13 @@ from dataclasses import replace from typing import Any, cast -from unittest.mock import AsyncMock, Mock +from unittest.mock import AsyncMock, Mock, patch import pytest from smithy_core.aio.client import AsyncClient from smithy_core.aio.eventstream import DuplexEventStream, InputEventStream from smithy_core.aio.interfaces import ClientTransport +from smithy_core.aio.retries import AdaptiveRetryStrategy from smithy_core.auth import AuthOption from smithy_core.exceptions import CallError, UnsupportedTransportError from smithy_core.response import EMPTY_RESPONSE_METADATA @@ -102,6 +103,37 @@ async def raise_modeled_error(**kwargs: Any) -> StubOutput: assert exc_info.value.response_metadata.http_status_code == 200 +async def test_response_metadata_kept_when_retry_times_out_waiting_for_send_token() -> ( + None +): + # Adaptive mode: the first attempt is throttled, then the retry can't get a send + # token. The caller should still get the throttle with its response metadata. + harness = pipeline_harness(NonDuplexTransport()) + + async def throttle(**kwargs: Any) -> StubOutput: + raise CallError("Rate exceeded", is_retry_safe=True, is_throttling_error=True) + + harness.protocol.deserialize_response = throttle # type: ignore[method-assign] + strategy = AdaptiveRetryStrategy(max_attempts=3) + timeout = TimeoutError("Failed to acquire 1.0 tokens within 30.0s") + + with ( + patch.object( + strategy, + "acquire_from_token_bucket", + AsyncMock(side_effect=[None, timeout]), + ), + patch("smithy_core.aio.client.sleep", AsyncMock()), + ): + with pytest.raises(CallError, match="Rate exceeded") as exc_info: + await harness.pipeline(retryable_client_call(strategy)) + + assert harness.transport.send_calls == 1 + assert exc_info.value.response_metadata.request_id == "stub-request-id" + assert exc_info.value.response_metadata.http_status_code == 200 + assert exc_info.value.__notes__ == [f"Retry not attempted: {timeout}"] + + async def test_response_metadata_empty_when_no_response_received() -> None: harness = pipeline_harness(NonDuplexTransport()) diff --git a/packages/smithy-core/tests/unit/aio/test_retries.py b/packages/smithy-core/tests/unit/aio/test_retries.py index ce7059bb9..c3425ff62 100644 --- a/packages/smithy-core/tests/unit/aio/test_retries.py +++ b/packages/smithy-core/tests/unit/aio/test_retries.py @@ -1,7 +1,11 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import gc +import weakref + import pytest from smithy_core.aio.retries import ( + AdaptiveRetryStrategy, RetryStrategyResolver, SimpleRetryStrategy, StandardRetryStrategy, @@ -286,6 +290,28 @@ async def test_retry_strategy_resolver_caches_strategies() -> None: assert strategy1 is not strategy3 +async def test_retry_strategy_resolver_cache_is_not_shared_across_resolvers() -> None: + resolver = RetryStrategyResolver() + options = RetryStrategyOptions(retry_mode="adaptive") + strategy = await resolver.resolve_retry_strategy(retry_strategy=options) + + for _ in range(200): + await RetryStrategyResolver().resolve_retry_strategy(retry_strategy=options) + + assert await resolver.resolve_retry_strategy(retry_strategy=options) is strategy + + +async def test_retry_strategy_resolver_strategies_are_freed_with_resolver() -> None: + resolver = RetryStrategyResolver() + strategy = await resolver.resolve_retry_strategy(retry_strategy=None) + strategy_ref = weakref.ref(strategy) + + del resolver, strategy + gc.collect() + + assert strategy_ref() is None + + async def test_retry_strategy_resolver_returns_existing_strategy() -> None: resolver = RetryStrategyResolver() provided_strategy = SimpleRetryStrategy(max_attempts=7) @@ -328,6 +354,17 @@ async def test_retry_strategy_resolver_uses_retry_mode_fallback() -> None: assert strategy.max_attempts == 4 +async def test_retry_strategy_resolver_uses_adaptive_retry_mode_fallback() -> None: + resolver = RetryStrategyResolver() + + strategy = await resolver.resolve_retry_strategy( + retry_strategy=None, retry_mode="adaptive", max_attempts=4 + ) + + assert isinstance(strategy, AdaptiveRetryStrategy) + assert strategy.max_attempts == 4 + + async def test_retry_strategy_resolver_fallback_defaults_when_unset() -> None: """Omitting both fallbacks must match the prior no-argument behavior.""" resolver = RetryStrategyResolver() diff --git a/packages/smithy-core/tests/unit/test_retries.py b/packages/smithy-core/tests/unit/test_retries.py index 4dc32d221..d81c8d687 100644 --- a/packages/smithy-core/tests/unit/test_retries.py +++ b/packages/smithy-core/tests/unit/test_retries.py @@ -1,12 +1,32 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import asyncio +import sys +import threading +import time +from collections.abc import Callable, Coroutine, Iterator +from copy import deepcopy +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + import pytest +from smithy_core.aio.client import RequestPipeline +from smithy_core.aio.retries import ( + AdaptiveRetryStrategy, + ClientRateLimiter, + RequestRateTracker, + RetryStrategyResolver, + TokenBucket, +) from smithy_core.exceptions import CallError, RetryError -from smithy_core.retries import ExponentialBackoffJitterType as EBJT from smithy_core.retries import ( + CubicCalculator, ExponentialRetryBackoffStrategy, + RetryStrategyOptions, StandardRetryQuota, ) +from smithy_core.retries import ExponentialBackoffJitterType as EBJT @pytest.mark.parametrize( @@ -126,3 +146,661 @@ def test_retry_quota_acquire_throttling_error( assert ( retry_quota.available_capacity == 28 - StandardRetryQuota.THROTTLING_RETRY_COST ) + + +class TestTokenBucket: + @pytest.mark.asyncio + async def test_initial_state(self): + token_bucket = TokenBucket() + assert token_bucket.current_capacity == 0 + assert token_bucket.max_capacity == token_bucket.MIN_CAPACITY + assert token_bucket.fill_rate == token_bucket.MIN_FILL_RATE + + @pytest.mark.asyncio + async def test_acquire_succeeds_immediately_within_capacity(self): + token_bucket = TokenBucket(curr_capacity=1) + + with patch("smithy_core.aio.retries.asyncio.wait_for") as mock_wait_for: + await token_bucket.acquire(1) + mock_wait_for.assert_not_called() + + assert token_bucket.current_capacity == 0 + + @pytest.mark.asyncio + async def test_acquire_waits_when_capacity_insufficient(self): + token_bucket = TokenBucket(fill_rate=100.0) + await token_bucket.acquire(1) + + with patch( + "smithy_core.aio.retries.asyncio.wait_for", wraps=asyncio.wait_for + ) as mock_wait_for: + await token_bucket.acquire(1) + mock_wait_for.assert_called() + + assert token_bucket.current_capacity == 0.0 + + @pytest.mark.asyncio + async def test_update_bucket_updates_rate(self): + token_bucket = TokenBucket(curr_capacity=1) + + await token_bucket.update_rate(5.0) + assert token_bucket.fill_rate == 5.0 + assert token_bucket.max_capacity == 5.0 + assert token_bucket.current_capacity == 1.0 + + @pytest.mark.asyncio + async def test_rate_can_never_be_zero(self): + token_bucket = TokenBucket() + await token_bucket.update_rate(0.0) + + assert token_bucket.fill_rate == TokenBucket.MIN_FILL_RATE + assert token_bucket.max_capacity == TokenBucket.MIN_CAPACITY + + @pytest.mark.asyncio + async def test_refill_caps_at_max_capacity(self): + # 100s at 4 tokens/s would refill 400 tokens, far more than the capacity of 4 + with patch("time.monotonic", side_effect=[0.0, 0.0, 100.0]): + token_bucket = TokenBucket() + await token_bucket.update_rate(4.0) + with token_bucket._lock: # type: ignore + token_bucket._refill() # type: ignore + + assert token_bucket.current_capacity == 4.0 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "actions,expected_capacity", + [ + ([("acquire", 1)], 0.0), + ([("acquire", 1), ("update", 4)], 1.0), + ([("acquire", 1), ("update", 4), ("acquire", 1)], 3.0), + ([("acquire", 1), ("update", 4), ("acquire", 1), ("acquire", 1)], 3.0), + ], + ) + async def test_multiple_refills_over_time( + self, actions: list[tuple[str, int]], expected_capacity: float + ): + time_values = [0.0, 1.0, 1.0, 1.5, 1.5, 4.0, 4.0, 5.0] + + with patch("time.monotonic", side_effect=time_values): + token_bucket = TokenBucket(curr_capacity=0, fill_rate=2.0) + + for action, value in actions: + if action == "acquire": + await token_bucket.acquire(value) + elif action == "update": + await token_bucket.update_rate(value) + + assert token_bucket.current_capacity == expected_capacity + + async def test_update_rate_wakes_waiting_acquirers(self): + token_bucket = TokenBucket(fill_rate=0.01) + + acquired = False + + async def slow_acquire(): + nonlocal acquired + await token_bucket.acquire(2.0) + acquired = True + + task = asyncio.create_task(slow_acquire()) + + # Let slow_acquire() start waiting + await asyncio.sleep(0.05) + assert not acquired + + await token_bucket.update_rate(100.0) + await asyncio.sleep(0.1) + assert acquired + await task + + async def test_acquire_does_not_wait_past_timeout(self): + # The bucket starts empty, and at the minimum fill rate a token takes 2s to + # refill, far past the timeout + token_bucket = TokenBucket(timeout=0.05) + + start = time.monotonic() + with pytest.raises(TimeoutError): + await token_bucket.acquire(1) + assert time.monotonic() - start < 1 + + def test_acquire_works_across_event_loops(self): + # A client, and so its bucket, can outlive the event loop it was first used on + token_bucket = TokenBucket(fill_rate=20.0) + + async def contended_acquires(): + # The bucket starts empty, so both acquires have to wait + await asyncio.gather(token_bucket.acquire(1), token_bucket.acquire(1)) + + asyncio.run(contended_acquires()) + asyncio.run(contended_acquires()) + + +class TestSharedAcrossThreads: + """One client, and so one rate limiter, may be used from several threads at once, + each running its own event loop.""" + + _THREADS = 8 + + @pytest.fixture + def frequent_thread_switches(self) -> Iterator[None]: + # Switch threads constantly so they collide inside the critical sections + original = sys.getswitchinterval() + sys.setswitchinterval(1e-6) + yield + sys.setswitchinterval(original) + + def _run_in_threads( + self, target: Callable[[], Coroutine[Any, Any, None]] + ) -> list[BaseException]: + errors: list[BaseException] = [] + + def run() -> None: + try: + asyncio.run(target()) + except BaseException as e: + errors.append(e) + + threads = [ + threading.Thread(target=run, daemon=True) for _ in range(self._THREADS) + ] + for thread in threads: + thread.start() + for thread in threads: + # A deadlock fails the test instead of hanging the suite + thread.join(timeout=30) + assert not thread.is_alive(), "thread did not finish" + return errors + + def test_rate_tracker(self, frequent_thread_switches: None): + tracker = RequestRateTracker() + + async def measure_many() -> None: + for _ in range(3000): + await tracker.measure_rate() + + assert self._run_in_threads(measure_many) == [] + + def test_token_bucket_rate_updates_and_acquires( + self, frequent_thread_switches: None + ): + bucket = TokenBucket(fill_rate=100000.0) + + async def update_and_acquire() -> None: + await bucket.update_rate(100000.0) + for _ in range(500): + await bucket.acquire(1) + assert bucket.current_capacity >= 0 + + assert self._run_in_threads(update_and_acquire) == [] + + def test_token_bucket_hands_out_each_token_once( + self, frequent_thread_switches: None + ): + acquires_per_thread = 500 + # Exactly enough tokens for every acquire, so no thread has to wait + bucket = TokenBucket(curr_capacity=self._THREADS * acquires_per_thread) + + async def spend() -> None: + for _ in range(acquires_per_thread): + await bucket.acquire(1) + + assert self._run_in_threads(spend) == [] + # A lost update to the count would leave tokens behind + assert bucket.current_capacity < 1 + + +class TestRateLimiter: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "rate_limiter_enabled, current_bucket_capacity, expected_bucket_capacity_after_acquire", + [ + (True, 2.0, 1.0), + (False, 1.0, 1.0), + ], + ) + async def test_token_consumption_before_request_scenarios( + self, + rate_limiter_enabled: bool, + current_bucket_capacity: float, + expected_bucket_capacity_after_acquire: float, + ): + token_bucket = TokenBucket(curr_capacity=current_bucket_capacity) + calculator = CubicCalculator(starting_max_rate=1.0, start_time=0.0) + tracker = RequestRateTracker() + + limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=calculator, + rate_tracker=tracker, + rate_limiter_enabled=rate_limiter_enabled, + ) + + await limiter.before_sending_request() + assert token_bucket.current_capacity == expected_bucket_capacity_after_acquire + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "throttling_error, rate_limiter_enabled, measured_rate, expected_bucket_fill_rate", + [ + (True, False, 5.0, 3.5), + (True, True, 5.0, 0.5), + (False, True, 5.0, 7.83), + (False, False, 5.0, 7.83), + ], + ) + async def test_token_consumption_after_response_scenarios( + self, + throttling_error: bool, + rate_limiter_enabled: bool, + measured_rate: float, + expected_bucket_fill_rate: float, + ): + with patch("time.monotonic") as mock_time: + mock_time.side_effect = [0.0, 0.1, 0.2, 0.3] + token_bucket = TokenBucket() + calculator = CubicCalculator(starting_max_rate=10.0, start_time=0.0) + tracker = RequestRateTracker() + + limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=calculator, + rate_tracker=tracker, + rate_limiter_enabled=rate_limiter_enabled, + ) + + with patch.object(tracker, "measure_rate", return_value=measured_rate): + await limiter.after_receiving_response( + throttling_error=throttling_error + ) + + assert round(token_bucket.fill_rate, 2) == expected_bucket_fill_rate + + async def test_throttling_response_enables_rate_limiter(self): + with patch("time.monotonic") as mock_time: + mock_time.side_effect = [0.0, 0.1, 0.2, 0.3, 0.4] + token_bucket = TokenBucket() + calculator = CubicCalculator(starting_max_rate=10.0, start_time=0.0) + tracker = RequestRateTracker() + + limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=calculator, + rate_tracker=tracker, + rate_limiter_enabled=False, + ) + + await limiter.after_receiving_response(throttling_error=True) + assert limiter.rate_limit_enabled is True + + async def test_raises_timeout_error_when_token_acquisition_takes_30_secs(self): + with patch("time.monotonic") as mock_time: + mock_time.side_effect = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 31] + token_bucket = TokenBucket(curr_capacity=1) + calculator = CubicCalculator() + tracker = RequestRateTracker() + limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=calculator, + rate_tracker=tracker, + rate_limiter_enabled=True, + ) + + await limiter.before_sending_request() + with pytest.raises(TimeoutError): + await limiter.before_sending_request() + + async def test_tokens_accumulation_after_numerous_successful_responses(self): + with patch("time.monotonic") as mock_time: + mock_time.side_effect = [i * 0.01 for i in range(5000)] + token_bucket = TokenBucket(curr_capacity=1) + calculator = CubicCalculator() + tracker = RequestRateTracker() + limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=calculator, + rate_tracker=tracker, + rate_limiter_enabled=True, + ) + + await limiter.after_receiving_response(throttling_error=True) + for _ in range(1000): + await limiter.after_receiving_response(throttling_error=False) + await limiter.after_receiving_response(throttling_error=True) + + # The final throttle caps the tokens saved up during the successes + assert token_bucket.current_capacity < 25 + + async def test_sending_rate_sequence(self): + # Each step is (throttled, timestamp, expected measured rate, expected fill + # rate of the bucket) + steps = [ + (False, 0.2, 0.0, 0.5), + (False, 0.4, 0.0, 0.5), + (False, 0.6, 4.8, 0.5), + (False, 0.8, 4.8, 0.5), + (False, 1.0, 4.16, 0.5), + (False, 1.2, 4.16, 0.6912), + (False, 1.4, 4.16, 1.0976), + (False, 1.6, 5.632, 1.6384), + (False, 1.8, 5.632, 2.3328), + (True, 2.0, 4.3264, 3.02848), + (False, 2.2, 4.3264, 3.486639), + (False, 2.4, 4.3264, 3.821874), + (False, 2.6, 5.66528, 4.053386), + (False, 2.8, 5.66528, 4.200373), + (False, 3.0, 4.333056, 4.282037), + (True, 3.2, 4.333056, 2.997426), + (False, 3.4, 4.333056, 3.452226), + ] + now = 0.0 + with patch( + "smithy_core.aio.retries.time", SimpleNamespace(monotonic=lambda: now) + ): + token_bucket = TokenBucket() + tracker = RequestRateTracker() + limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=CubicCalculator(start_time=0.0), + rate_tracker=tracker, + rate_limiter_enabled=False, + ) + + for throttled, timestamp, measured_rate, fill_rate in steps: + now = timestamp + await limiter.after_receiving_response(throttling_error=throttled) + + assert tracker.measured_client_rate == pytest.approx( # type: ignore + measured_rate, abs=1e-6 + ) + assert token_bucket.fill_rate == pytest.approx(fill_rate, abs=1e-6) # type: ignore + + async def test_calculated_rate_is_capped_at_2x_measured_rate(self): + with patch("time.monotonic") as mock_time: + mock_time.side_effect = [0.0, 0.1, 0.2, 0.3, 0.4] + token_bucket = TokenBucket() + calculator = CubicCalculator(starting_max_rate=10.0, start_time=0.0) + tracker = RequestRateTracker() + + limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=calculator, + rate_tracker=tracker, + rate_limiter_enabled=True, + ) + + with ( + patch.object(tracker, "measure_rate", return_value=5.0), + patch.object( + calculator, "calculate_scaled_request_rate", return_value=20.0 + ), + patch.object(token_bucket, "_set_rate") as mock_set_rate, + ): + await limiter.after_receiving_response(throttling_error=False) + + # calculated rate should be capped at 2x measured rate at 2 * 5.0 = 10.0 + mock_set_rate.assert_called_once_with(10.0) + + async def test_new_rate_is_set_while_holding_the_lock(self): + # Two threads could otherwise calculate rates A then B but apply B then A + token_bucket = TokenBucket() + limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=CubicCalculator(), + rate_tracker=RequestRateTracker(), + rate_limiter_enabled=True, + ) + lock_held: list[bool] = [] + + def record_lock(rate: float) -> None: + lock_held.append(limiter._lock.locked()) # type: ignore[reportPrivateUsage] + + token_bucket._set_rate = record_lock # type: ignore[method-assign] + await limiter.after_receiving_response(throttling_error=True) + await limiter.after_receiving_response(throttling_error=False) + + assert lock_held == [True, True] + + +class TestCubicCalculator: + def test_cubic_calculator_initializes(self): + calculator = CubicCalculator(starting_max_rate=10.0, start_time=5.0) + assert calculator.last_max_rate == 10.0 + assert calculator.last_throttle_time == 5.0 + + def test_calculate_inflection_point(self): + calculator = CubicCalculator(starting_max_rate=10.0, start_time=5.0) + time_window = calculator.calculate_and_update_inflection_point() + assert round(time_window, 2) == 1.96 + + @pytest.mark.parametrize( + "cases", + [ + # (response, timestamp, calculated_rate) + [ + ("success", 5, 7.0), + ("success", 6, 9.64893600966), + ("success", 7, 10.000030849917364), + ("success", 8, 10.453284520772092), + ("success", 9, 13.408697022224185), + ("success", 10, 21.26626835427364), + ("success", 11, 36.425998516920465), + ], + [ + ("success", 5, 7.0), + ("success", 6, 9.64893600966), + ("throttle", 7, 6.754255206761999), + ("throttle", 8, 4.727978644733399), + ("success", 9, 6.606547753887045), + ("success", 10, 6.763279816944947), + ("success", 11, 7.598174833907107), + ("success", 12, 11.511232804773524), + ], + ], + ids=["successes", "successes_and_throttles"], + ) + def test_rate_sequence_after_successes_and_throttles( + self, cases: list[tuple[str, float, float]] + ): + calculator = CubicCalculator(starting_max_rate=10.0, start_time=5.0) + current_rate = 10.0 + for response, timestamp, expected in cases: + if response == "success": + current_rate = calculator.calculate_scaled_request_rate(timestamp) + else: + current_rate = calculator.calculate_throttled_request_rate( + current_rate, timestamp + ) + assert current_rate == pytest.approx(expected) # type: ignore + + def test_throttle_request(self): + calculator = CubicCalculator(starting_max_rate=10.0, start_time=5.0) + result = calculator.calculate_throttled_request_rate( + rate_to_use=8.0, timestamp=10.0 + ) + + assert result == 8.0 * 0.7 + assert calculator.last_max_rate == 8.0 + assert calculator.last_throttle_time == 10.0 + + +class TestRequestRateTracker: + async def test_measure_rate_same_bucket(self): + with patch("time.monotonic") as mock_time: + mock_time.side_effect = [0.0, 0.3, 0.4] + tracker = RequestRateTracker() + rate1 = await tracker.measure_rate() + rate2 = await tracker.measure_rate() + + assert rate1 == 0 + assert rate2 == 0 + assert tracker.request_count == 2 + + async def test_measure_rate_new_bucket(self): + with patch("time.monotonic") as mock_time: + mock_time.side_effect = [0.0, 0.1, 0.7] + tracker = RequestRateTracker() + + await tracker.measure_rate() + rate = await tracker.measure_rate() + + # 2 requests over a 0.5s bucket is 4 r/s, smoothed: 4 * 0.8 + 0 * 0.2 + assert rate == pytest.approx(3.2) # type: ignore + assert tracker.request_count == 0 + + +class TestAdaptiveRetryStrategy: + """Tests for AdaptiveRetryStrategy with rate limiting.""" + + def test_initialization_with_default_rate_limiter(self): + strategy = AdaptiveRetryStrategy() + + assert strategy.rate_limiter is not None + assert isinstance(strategy.rate_limiter, ClientRateLimiter) + assert strategy.rate_limiter.rate_limit_enabled is False + + def test_initialization_with_custom_rate_limiter(self): + token_bucket = TokenBucket() + calculator = CubicCalculator(starting_max_rate=5.0, start_time=0.0) + tracker = RequestRateTracker() + custom_limiter = ClientRateLimiter( + token_bucket=token_bucket, + cubic_calculator=calculator, + rate_tracker=tracker, + rate_limiter_enabled=True, + ) + + strategy = AdaptiveRetryStrategy(rate_limiter=custom_limiter) + + assert strategy.rate_limiter is custom_limiter + assert strategy.rate_limiter.rate_limit_enabled is True + + @pytest.mark.parametrize( + "error, expected", + [ + (CallError(message="Throttled", is_throttling_error=True), True), + (CallError(message="Server error", is_throttling_error=False), False), + (Exception("Generic error"), False), + ], + ids=["throttling", "not_throttling", "not_a_call_error"], + ) + def test_is_throttling_error(self, error: Exception, expected: bool): + assert AdaptiveRetryStrategy().is_throttling_error(error) is expected + + @pytest.mark.asyncio + async def test_resolver_creates_adaptive_strategy(self): + resolver = RetryStrategyResolver() + option1 = RetryStrategyOptions(retry_mode="adaptive") + + strategy = await resolver.resolve_retry_strategy(retry_strategy=option1) + + assert isinstance(strategy, AdaptiveRetryStrategy) + assert strategy.max_attempts == 3 + + @pytest.mark.parametrize("max_attempts", [2, 3, 10]) + async def test_adaptive_retry_strategy(self, max_attempts: int) -> None: + strategy = AdaptiveRetryStrategy(max_attempts=max_attempts) + error = CallError(is_retry_safe=True) + token = await strategy.acquire_initial_retry_token() + for _ in range(max_attempts - 1): + token = await strategy.refresh_retry_token_for_retry( + token_to_renew=token, error=error + ) + with pytest.raises(RetryError): + await strategy.refresh_retry_token_for_retry( + token_to_renew=token, error=error + ) + + def test_deepcopy_returns_adaptive_strategy_with_rate_limiter(self) -> None: + strategy = AdaptiveRetryStrategy() + copied = deepcopy(strategy) + + assert copied is strategy + assert isinstance(copied, AdaptiveRetryStrategy) + assert copied.rate_limiter is strategy.rate_limiter + + +class TestRequestPipelineRateLimiting: + def _non_retryable_call(self, strategy: AdaptiveRetryStrategy) -> Any: + # A call whose input stream can't be resent, so it skips the retry loop + call = MagicMock() + call.retryable.return_value = False + call.retry_strategy = strategy + return call + + def _adaptive_strategy_with_mocked_limiter( + self, + ) -> tuple[AdaptiveRetryStrategy, ClientRateLimiter]: + limiter = ClientRateLimiter( + token_bucket=TokenBucket(), + cubic_calculator=CubicCalculator(), + rate_tracker=RequestRateTracker(), + rate_limiter_enabled=True, + ) + limiter.before_sending_request = AsyncMock() # type: ignore[method-assign] + limiter.after_receiving_response = AsyncMock() # type: ignore[method-assign] + return AdaptiveRetryStrategy(rate_limiter=limiter), limiter + + @pytest.mark.asyncio + async def test_non_retryable_call_applies_rate_limiting_on_success(self): + strategy, limiter = self._adaptive_strategy_with_mocked_limiter() + pipeline: RequestPipeline[Any, Any] = RequestPipeline( + protocol=MagicMock(), transport=MagicMock() + ) + handle_attempt = AsyncMock( + return_value=SimpleNamespace(transport_response=MagicMock(), response="ok") + ) + + with patch.object(pipeline, "_handle_attempt", handle_attempt): + await pipeline._retry( # type: ignore[reportPrivateUsage] + self._non_retryable_call(strategy), MagicMock(), None + ) + + handle_attempt.assert_awaited_once() + limiter.before_sending_request.assert_awaited_once() # type: ignore[attr-defined] + limiter.after_receiving_response.assert_awaited_once_with( # type: ignore[attr-defined] + throttling_error=False + ) + + @pytest.mark.asyncio + async def test_non_retryable_call_reports_throttle_without_retrying(self): + strategy, limiter = self._adaptive_strategy_with_mocked_limiter() + pipeline: RequestPipeline[Any, Any] = RequestPipeline( + protocol=MagicMock(), transport=MagicMock() + ) + error = CallError(is_retry_safe=True, is_throttling_error=True) + handle_attempt = AsyncMock( + return_value=SimpleNamespace(transport_response=MagicMock(), response=error) + ) + + with patch.object(pipeline, "_handle_attempt", handle_attempt): + output_context = await pipeline._retry( # type: ignore[reportPrivateUsage] + self._non_retryable_call(strategy), MagicMock(), None + ) + + # Sent once and returned as-is, even though the error is retry-safe + handle_attempt.assert_awaited_once() + assert output_context.response is error + limiter.before_sending_request.assert_awaited_once() # type: ignore[attr-defined] + limiter.after_receiving_response.assert_awaited_once_with(True) # type: ignore[attr-defined] + + @pytest.mark.asyncio + async def test_failure_without_response_does_not_update_rate_limiter(self): + strategy, limiter = self._adaptive_strategy_with_mocked_limiter() + pipeline: RequestPipeline[Any, Any] = RequestPipeline( + protocol=MagicMock(), transport=MagicMock() + ) + # Fails before sending, e.g. endpoint resolution or signing, so no response + error = CallError(message="Unable to resolve endpoint", is_retry_safe=False) + handle_attempt = AsyncMock( + return_value=SimpleNamespace(transport_response=None, response=error) + ) + + with patch.object(pipeline, "_handle_attempt", handle_attempt): + output_context = await pipeline._retry( # type: ignore[reportPrivateUsage] + self._non_retryable_call(strategy), MagicMock(), None + ) + + assert output_context.response is error + limiter.before_sending_request.assert_awaited_once() # type: ignore[attr-defined] + limiter.after_receiving_response.assert_not_awaited() # type: ignore[attr-defined]