From 7e0ef1ce240e19cf66a37885f1f2a8edd630fde0 Mon Sep 17 00:00:00 2001 From: SamRemis Date: Fri, 7 Nov 2025 08:05:32 -0500 Subject: [PATCH 01/11] Working commit --- .../python/codegen/ClientGenerator.java | 4 +- .../codegen/generators/ConfigGenerator.java | 27 +++++++++---- .../smithy-core/src/smithy_core/aio/client.py | 8 ++-- .../src/smithy_core/interfaces/retries.py | 17 +++++++- .../smithy-core/src/smithy_core/retries.py | 39 ++++++++++++++++++- 5 files changed, 81 insertions(+), 14 deletions(-) diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index 1cac89621..2295836e7 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -207,12 +207,12 @@ raise ExpectationNotMetError("protocol and transport MUST be set on the config t call = ClientCall( input=input, operation=${operation:T}, - context=TypedProperties({"config": config}), + context=TypedProperties({"config": config, "client_object_id": str(id(self))}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, - retry_strategy=config.retry_strategy, + retry_strategy_resolver=config.retry_strategy_resolver, ) """, writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins))); diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java index 45b6324d7..579d38bc4 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java @@ -53,18 +53,31 @@ public final class ConfigGenerator implements Runnable { .initialize(writer -> writer.write("self.interceptors = interceptors or []")) .build(), ConfigProperty.builder() - .name("retry_strategy") + .name("retry_strategy_resolver") .type(Symbol.builder() - .name("RetryStrategy") - .namespace("smithy_core.interfaces.retries", ".") - .addDependency(SmithyPythonDependency.SMITHY_CORE) + .name("RetryStrategyResolver[RetryStrategy, Mapping[str, Any]]") + .addReference(Symbol.builder() + .name("RetryStrategyResolver") + .namespace("smithy_core.interfaces.retries", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) + .addReference(Symbol.builder() + .name("RetryStrategy") + .namespace("smithy_core.interfaces.retries", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) + .addReference(Symbol.builder() + .name("Mapping") + .namespace("collections.abc", ".") + .putProperty(SymbolProperties.STDLIB, true) + .build()) .build()) - .documentation("The retry strategy for issuing retry tokens and computing retry delays.") + .documentation("The retry strategy resolver for resolving retry strategies per client.") .nullable(false) .initialize(writer -> { writer.addDependency(SmithyPythonDependency.SMITHY_CORE); - writer.addImport("smithy_core.retries", "SimpleRetryStrategy"); - writer.write("self.retry_strategy = retry_strategy or SimpleRetryStrategy()"); + writer.addImport("smithy_core.retries", "CachingRetryStrategyResolver"); + writer.write("self.retry_strategy_resolver = retry_strategy_resolver or CachingRetryStrategyResolver()"); }) .build(), ConfigProperty.builder() diff --git a/packages/smithy-core/src/smithy_core/aio/client.py b/packages/smithy-core/src/smithy_core/aio/client.py index bf27c440c..5a990a912 100644 --- a/packages/smithy-core/src/smithy_core/aio/client.py +++ b/packages/smithy-core/src/smithy_core/aio/client.py @@ -22,7 +22,7 @@ ) from ..interfaces import Endpoint, TypedProperties from ..interfaces.auth import AuthOption, AuthSchemeResolver -from ..interfaces.retries import RetryStrategy +from ..interfaces.retries import RetryStrategyResolver from ..schemas import APIOperation from ..serializers import SerializeableShape from ..shapes import ShapeID @@ -77,7 +77,7 @@ class ClientCall[I: SerializeableShape, O: DeserializeableShape]: endpoint_resolver: EndpointResolver """The endpoint resolver for the operation.""" - retry_strategy: RetryStrategy + retry_strategy_resolver: RetryStrategyResolver """The retry strategy to use for the operation.""" retry_scope: str | None = None @@ -329,7 +329,9 @@ async def _retry[I: SerializeableShape, O: DeserializeableShape]( if not call.retryable(): return await self._handle_attempt(call, request_context, request_future) - retry_strategy = call.retry_strategy + retry_strategy = await call.retry_strategy_resolver.get_retry_strategy( + properties=request_context.properties + ) retry_token = retry_strategy.acquire_initial_retry_token( token_scope=call.retry_scope ) diff --git a/packages/smithy-core/src/smithy_core/interfaces/retries.py b/packages/smithy-core/src/smithy_core/interfaces/retries.py index a5c9d428b..886e21be0 100644 --- a/packages/smithy-core/src/smithy_core/interfaces/retries.py +++ b/packages/smithy-core/src/smithy_core/interfaces/retries.py @@ -1,7 +1,8 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Mapping from dataclasses import dataclass -from typing import Protocol, runtime_checkable +from typing import Any, Literal, Protocol, runtime_checkable @runtime_checkable @@ -52,6 +53,9 @@ class RetryToken(Protocol): """Delay in seconds to wait before the retry attempt.""" +RetryStrategyType = Literal["simple"] + + class RetryStrategy(Protocol): """Issuer of :py:class:`RetryToken`s.""" @@ -100,3 +104,14 @@ def record_success(self, *, token: RetryToken) -> None: :param token: The token used for the previous successful attempt. """ ... + + +class RetryStrategyResolver[RS: RetryStrategy, P: Mapping[str, Any]](Protocol): + """Used to resolve a RetryStrategy for a given caller.""" + + async def get_retry_strategy(self, *, properties: P) -> RS: + """Get the retry strategy for the caller. + + :param properties: Properties including caller identification. + """ + ... diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index 06bf6f988..e21476670 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -1,12 +1,49 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import asyncio import random -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from enum import Enum +from functools import lru_cache +from typing import Any from .exceptions import RetryError from .interfaces import retries as retries_interface +from .interfaces.retries import RetryStrategy, RetryStrategyResolver, RetryStrategyType + + +class CachingRetryStrategyResolver[RS: RetryStrategy, P: Mapping[str, Any]]( + RetryStrategyResolver[RS, P] +): + def __init__(self) -> None: + self._locks: dict[str, asyncio.Lock] = {} + self._main_lock = asyncio.Lock() + + @lru_cache(maxsize=50) + def _create_retry_strategy_cached( + self, caller_id: str, retry_strategy: RetryStrategyType + ) -> RS: + return self._create_retry_strategy(retry_strategy) + + async def get_retry_strategy(self, *, properties: P) -> RS: + caller_id = properties["client_object_id"] + retry_strategy = properties.get("retry_strategy", "simple") + + async with self._main_lock: + if caller_id not in self._locks: + self._locks[caller_id] = asyncio.Lock() + lock = self._locks[caller_id] + + async with lock: + return self._create_retry_strategy_cached(caller_id, retry_strategy) + + def _create_retry_strategy(self, retry_strategy: RetryStrategyType) -> RS: + match retry_strategy: + case "simple": + return SimpleRetryStrategy() + case _: + raise ValueError(f"Unknown retry strategy: {retry_strategy}") class ExponentialBackoffJitterType(Enum): From 1b6c10c95f9403fbd54ba91fa5d2405f752ebb3f Mon Sep 17 00:00:00 2001 From: SamRemis Date: Fri, 7 Nov 2025 12:14:10 -0500 Subject: [PATCH 02/11] First revision --- .../python/codegen/ClientGenerator.java | 3 +- .../smithy-core/src/smithy_core/aio/client.py | 16 +++-- .../src/smithy_core/interfaces/retries.py | 13 ++-- .../smithy-core/src/smithy_core/retries.py | 62 ++++++++++++---- .../smithy-core/tests/unit/test_retries.py | 70 ++++++++++++++++++- 5 files changed, 138 insertions(+), 26 deletions(-) diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index 2295836e7..3090e760e 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -183,6 +183,7 @@ private void writeSharedOperationInit(PythonWriter writer, OperationShape operat writer.putContext("operation", symbolProvider.toSymbol(operation)); writer.addImport("smithy_core.aio.client", "ClientCall"); + writer.addImport("smithy_core.aio.client", "CLIENT_ID"); writer.addImport("smithy_core.interceptors", "InterceptorChain"); writer.addImport("smithy_core.types", "TypedProperties"); writer.addImport("smithy_core.aio.client", "RequestPipeline"); @@ -207,7 +208,7 @@ raise ExpectationNotMetError("protocol and transport MUST be set on the config t call = ClientCall( input=input, operation=${operation:T}, - context=TypedProperties({"config": config, "client_object_id": str(id(self))}), + context=TypedProperties({"config": config, CLIENT_ID: str(id(self))}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, diff --git a/packages/smithy-core/src/smithy_core/aio/client.py b/packages/smithy-core/src/smithy_core/aio/client.py index 5a990a912..b5db60ec8 100644 --- a/packages/smithy-core/src/smithy_core/aio/client.py +++ b/packages/smithy-core/src/smithy_core/aio/client.py @@ -22,7 +22,7 @@ ) from ..interfaces import Endpoint, TypedProperties from ..interfaces.auth import AuthOption, AuthSchemeResolver -from ..interfaces.retries import RetryStrategyResolver +from ..interfaces.retries import RetryStrategy, RetryStrategyResolver from ..schemas import APIOperation from ..serializers import SerializeableShape from ..shapes import ShapeID @@ -44,6 +44,14 @@ AUTH_SCHEME = PropertyKey(key="auth_scheme", value_type=AuthScheme[Any, Any, Any, Any]) +CLIENT_ID = PropertyKey(key="client_id", value_type=str) +"""A unique identifier for the client instance. + +This property key is used to identify a specific client instance across operations. +It enables per-client caching and state management in resolvers and other components. +The value is typically the string representation of the client's object ID. +""" + _UNRESOLVED = URI(host="", path="/") _LOGGER = logging.getLogger(__name__) @@ -77,8 +85,8 @@ class ClientCall[I: SerializeableShape, O: DeserializeableShape]: endpoint_resolver: EndpointResolver """The endpoint resolver for the operation.""" - retry_strategy_resolver: RetryStrategyResolver - """The retry strategy to use for the operation.""" + retry_strategy_resolver: RetryStrategyResolver[RetryStrategy] + """The retry strategy resolver for the operation.""" retry_scope: str | None = None """The retry scope for the operation.""" @@ -329,7 +337,7 @@ async def _retry[I: SerializeableShape, O: DeserializeableShape]( if not call.retryable(): return await self._handle_attempt(call, request_context, request_future) - retry_strategy = await call.retry_strategy_resolver.get_retry_strategy( + retry_strategy = await call.retry_strategy_resolver.resolve_retry_strategy( properties=request_context.properties ) retry_token = retry_strategy.acquire_initial_retry_token( diff --git a/packages/smithy-core/src/smithy_core/interfaces/retries.py b/packages/smithy-core/src/smithy_core/interfaces/retries.py index 886e21be0..374d10fd2 100644 --- a/packages/smithy-core/src/smithy_core/interfaces/retries.py +++ b/packages/smithy-core/src/smithy_core/interfaces/retries.py @@ -1,8 +1,9 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Literal, Protocol, runtime_checkable +from typing import Literal, Protocol, runtime_checkable + +from . import TypedProperties @runtime_checkable @@ -106,12 +107,12 @@ def record_success(self, *, token: RetryToken) -> None: ... -class RetryStrategyResolver[RS: RetryStrategy, P: Mapping[str, Any]](Protocol): +class RetryStrategyResolver[RS: RetryStrategy](Protocol): """Used to resolve a RetryStrategy for a given caller.""" - async def get_retry_strategy(self, *, properties: P) -> RS: - """Get the retry strategy for the caller. + async def resolve_retry_strategy(self, *, properties: TypedProperties) -> RS: + """Resolve the retry strategy for the caller. - :param properties: Properties including caller identification. + :param properties: Properties including caller identification and config. """ ... diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index e21476670..97de28e0e 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -2,43 +2,79 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio import random -from collections.abc import Callable, Mapping +from collections.abc import Callable from dataclasses import dataclass from enum import Enum from functools import lru_cache from typing import Any +from .aio.client import CLIENT_ID from .exceptions import RetryError from .interfaces import retries as retries_interface from .interfaces.retries import RetryStrategy, RetryStrategyResolver, RetryStrategyType -class CachingRetryStrategyResolver[RS: RetryStrategy, P: Mapping[str, Any]]( - RetryStrategyResolver[RS, P] -): +class CachingRetryStrategyResolver(RetryStrategyResolver[RetryStrategy]): + """Caching retry strategy resolver that creates and caches retry strategies per caller. + + This resolver maintains a cache of retry strategies keyed by a unique identifier + for each caller. This allows multiple operations from the same caller to share + a single retry strategy instance, which is important for strategies that maintain + state across retries (e.g., token buckets, rate limiters). + + The resolver uses async locks to ensure thread-safe creation of strategies when + multiple concurrent requests arrive for the same caller. + """ + def __init__(self) -> None: self._locks: dict[str, asyncio.Lock] = {} self._main_lock = asyncio.Lock() + def __deepcopy__(self, memo: dict[int, Any]) -> "CachingRetryStrategyResolver": + """Return self to preserve cache across config copies. + + The resolver is designed to be shared across operations from the same client, + so it should not be deep copied when the config is copied per-operation. + This ensures the cache and locks are preserved. + """ + return self + @lru_cache(maxsize=50) def _create_retry_strategy_cached( - self, caller_id: str, retry_strategy: RetryStrategyType - ) -> RS: + self, retry_id: str, retry_strategy: RetryStrategyType + ) -> RetryStrategy: return self._create_retry_strategy(retry_strategy) - async def get_retry_strategy(self, *, properties: P) -> RS: - caller_id = properties["client_object_id"] + async def resolve_retry_strategy( + self, *, properties: retries_interface.TypedProperties + ) -> RetryStrategy: + """Get or create a retry strategy for the caller. + + :param properties: Properties map that must contain the CLIENT_ID property key + with a unique identifier for the caller. This ID is used to cache and retrieve + the appropriate retry strategy instance. Optionally may contain a + "retry_strategy" key to specify the strategy type (defaults to "simple"). + :raises ValueError: If CLIENT_ID is not present in properties. + """ + retry_id = properties.get(CLIENT_ID.key) + if retry_id is None: + raise ValueError( + f"Properties must contain '{CLIENT_ID.key}' key with a unique identifier for the caller" + ) + retry_strategy = properties.get("retry_strategy", "simple") async with self._main_lock: - if caller_id not in self._locks: - self._locks[caller_id] = asyncio.Lock() - lock = self._locks[caller_id] + if retry_id not in self._locks: + self._locks[retry_id] = asyncio.Lock() + lock = self._locks[retry_id] async with lock: - return self._create_retry_strategy_cached(caller_id, retry_strategy) + return self._create_retry_strategy_cached(retry_id, retry_strategy) - def _create_retry_strategy(self, retry_strategy: RetryStrategyType) -> RS: + def _create_retry_strategy( + self, retry_strategy: RetryStrategyType + ) -> RetryStrategy: match retry_strategy: case "simple": return SimpleRetryStrategy() diff --git a/packages/smithy-core/tests/unit/test_retries.py b/packages/smithy-core/tests/unit/test_retries.py index 0b3c23be4..963570ce9 100644 --- a/packages/smithy-core/tests/unit/test_retries.py +++ b/packages/smithy-core/tests/unit/test_retries.py @@ -1,10 +1,19 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import asyncio import pytest +from smithy_core.aio.client import CLIENT_ID from smithy_core.exceptions import CallError, RetryError -from smithy_core.retries import ExponentialBackoffJitterType as EBJT -from smithy_core.retries import ExponentialRetryBackoffStrategy, SimpleRetryStrategy +from smithy_core.retries import ( + CachingRetryStrategyResolver, + ExponentialRetryBackoffStrategy, + SimpleRetryStrategy, +) +from smithy_core.retries import ( + ExponentialBackoffJitterType as EBJT, +) +from smithy_core.types import TypedProperties @pytest.mark.parametrize( @@ -100,3 +109,60 @@ def test_simple_retry_does_not_retry_unsafe() -> None: token = strategy.acquire_initial_retry_token() with pytest.raises(RetryError): strategy.refresh_retry_token_for_retry(token_to_renew=token, error=error) + + +@pytest.mark.asyncio +async def test_caching_retry_strategy_default_resolution() -> None: + resolver = CachingRetryStrategyResolver() + properties = TypedProperties({CLIENT_ID.key: "test-client-1"}) + + strategy = await resolver.resolve_retry_strategy(properties=properties) + + assert isinstance(strategy, SimpleRetryStrategy) + + +@pytest.mark.asyncio +async def test_caching_retry_strategy_resolver_caches_per_client() -> None: + resolver = CachingRetryStrategyResolver() + properties1 = TypedProperties({CLIENT_ID.key: "test-caller-1"}) + properties2 = TypedProperties({CLIENT_ID.key: "test-caller-2"}) + + strategy1a = await resolver.resolve_retry_strategy(properties=properties1) + strategy1b = await resolver.resolve_retry_strategy(properties=properties1) + strategy2 = await resolver.resolve_retry_strategy(properties=properties2) + + assert strategy1a is strategy1b + assert strategy1a is not strategy2 + + +@pytest.mark.asyncio +async def test_caching_retry_strategy_resolver_concurrent_access() -> None: + resolver = CachingRetryStrategyResolver() + properties = TypedProperties({CLIENT_ID.key: "test-caller-concurrent"}) + + strategies = await asyncio.gather( + resolver.resolve_retry_strategy(properties=properties), + resolver.resolve_retry_strategy(properties=properties), + resolver.resolve_retry_strategy(properties=properties), + ) + + assert strategies[0] is strategies[1] + assert strategies[1] is strategies[2] + + +@pytest.mark.asyncio +async def test_caching_retry_strategy_resolver_requires_client_id() -> None: + resolver = CachingRetryStrategyResolver() + properties = TypedProperties({}) + + with pytest.raises(ValueError, match=CLIENT_ID.key): + await resolver.resolve_retry_strategy(properties=properties) + + +def test_caching_retry_strategy_resolver_survives_deepcopy() -> None: + from copy import deepcopy + + resolver = CachingRetryStrategyResolver() + resolver_copy = deepcopy(resolver) + + assert resolver is resolver_copy From 2c14d209a3865fc399c37df81633664fee1cf740 Mon Sep 17 00:00:00 2001 From: SamRemis Date: Fri, 7 Nov 2025 15:15:03 -0500 Subject: [PATCH 03/11] Add retry options --- .../codegen/generators/ConfigGenerator.java | 25 ++++++-- .../smithy-core/src/smithy_core/aio/client.py | 4 -- .../smithy-core/src/smithy_core/retries.py | 58 +++++++++++++------ .../smithy-core/tests/unit/test_retries.py | 36 ++++++++++-- 4 files changed, 90 insertions(+), 33 deletions(-) diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java index 579d38bc4..dd3842601 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java @@ -55,7 +55,7 @@ public final class ConfigGenerator implements Runnable { ConfigProperty.builder() .name("retry_strategy_resolver") .type(Symbol.builder() - .name("RetryStrategyResolver[RetryStrategy, Mapping[str, Any]]") + .name("RetryStrategyResolver[RetryStrategy]") .addReference(Symbol.builder() .name("RetryStrategyResolver") .namespace("smithy_core.interfaces.retries", ".") @@ -66,11 +66,6 @@ public final class ConfigGenerator implements Runnable { .namespace("smithy_core.interfaces.retries", ".") .addDependency(SmithyPythonDependency.SMITHY_CORE) .build()) - .addReference(Symbol.builder() - .name("Mapping") - .namespace("collections.abc", ".") - .putProperty(SymbolProperties.STDLIB, true) - .build()) .build()) .documentation("The retry strategy resolver for resolving retry strategies per client.") .nullable(false) @@ -80,6 +75,24 @@ public final class ConfigGenerator implements Runnable { writer.write("self.retry_strategy_resolver = retry_strategy_resolver or CachingRetryStrategyResolver()"); }) .build(), + ConfigProperty.builder() + .name("retry_options") + .type(Symbol.builder() + .name("RetryStrategyOptions") + .addReference(Symbol.builder() + .name("RetryStrategyOptions") + .namespace("smithy_core.retries", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) + .build()) + .documentation("Options for configuring retry behavior.") + .nullable(false) + .initialize(writer -> { + writer.addDependency(SmithyPythonDependency.SMITHY_CORE); + writer.addImport("smithy_core.retries", "RetryStrategyOptions"); + writer.write("self.retry_options = retry_options or RetryStrategyOptions()"); + }) + .build(), ConfigProperty.builder() .name("endpoint_uri") .type(Symbol.builder() diff --git a/packages/smithy-core/src/smithy_core/aio/client.py b/packages/smithy-core/src/smithy_core/aio/client.py index b5db60ec8..08a061588 100644 --- a/packages/smithy-core/src/smithy_core/aio/client.py +++ b/packages/smithy-core/src/smithy_core/aio/client.py @@ -46,10 +46,6 @@ CLIENT_ID = PropertyKey(key="client_id", value_type=str) """A unique identifier for the client instance. - -This property key is used to identify a specific client instance across operations. -It enables per-client caching and state management in resolvers and other components. -The value is typically the string representation of the client's object ID. """ _UNRESOLVED = URI(host="", path="/") diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index 97de28e0e..a846df635 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -6,12 +6,34 @@ from dataclasses import dataclass from enum import Enum from functools import lru_cache -from typing import Any +from typing import Any, Protocol from .aio.client import CLIENT_ID from .exceptions import RetryError from .interfaces import retries as retries_interface from .interfaces.retries import RetryStrategy, RetryStrategyResolver, RetryStrategyType +from .types import PropertyKey + + +@dataclass(kw_only=True, frozen=True) +class RetryStrategyOptions: + """Options for configuring retry behavior. + """ + + retry_mode: RetryStrategyType = "simple" + """The retry mode to use.""" + + max_attempts: int = 3 + """Maximum number of attempts (initial attempt plus retries).""" + + +class RetryConfig(Protocol): + """Protocol for config objects that support retry configuration.""" + + retry_options: RetryStrategyOptions + + +RETRY_CONFIG = PropertyKey(key="config", value_type=RetryConfig) class CachingRetryStrategyResolver(RetryStrategyResolver[RetryStrategy]): @@ -22,8 +44,6 @@ class CachingRetryStrategyResolver(RetryStrategyResolver[RetryStrategy]): a single retry strategy instance, which is important for strategies that maintain state across retries (e.g., token buckets, rate limiters). - The resolver uses async locks to ensure thread-safe creation of strategies when - multiple concurrent requests arrive for the same caller. """ def __init__(self) -> None: @@ -31,19 +51,15 @@ def __init__(self) -> None: self._main_lock = asyncio.Lock() def __deepcopy__(self, memo: dict[int, Any]) -> "CachingRetryStrategyResolver": - """Return self to preserve cache across config copies. - - The resolver is designed to be shared across operations from the same client, - so it should not be deep copied when the config is copied per-operation. - This ensures the cache and locks are preserved. + """Return self to preserve cache across operation-level config copies. """ return self @lru_cache(maxsize=50) def _create_retry_strategy_cached( - self, retry_id: str, retry_strategy: RetryStrategyType + self, retry_id: str, retry_mode: RetryStrategyType, max_attempts: int ) -> RetryStrategy: - return self._create_retry_strategy(retry_strategy) + return self._create_retry_strategy(retry_mode, max_attempts) async def resolve_retry_strategy( self, *, properties: retries_interface.TypedProperties @@ -51,9 +67,9 @@ async def resolve_retry_strategy( """Get or create a retry strategy for the caller. :param properties: Properties map that must contain the CLIENT_ID property key - with a unique identifier for the caller. This ID is used to cache and retrieve - the appropriate retry strategy instance. Optionally may contain a - "retry_strategy" key to specify the strategy type (defaults to "simple"). + with a unique identifier for the caller, and a "config" key with a + retry_strategy attribute (RetryStrategyOptions) specifying the strategy + configuration. Strategies are cached per client and options combination. :raises ValueError: If CLIENT_ID is not present in properties. """ retry_id = properties.get(CLIENT_ID.key) @@ -62,7 +78,9 @@ async def resolve_retry_strategy( f"Properties must contain '{CLIENT_ID.key}' key with a unique identifier for the caller" ) - retry_strategy = properties.get("retry_strategy", "simple") + # Get retry options from config + config = properties[RETRY_CONFIG] + options = config.retry_options async with self._main_lock: if retry_id not in self._locks: @@ -70,16 +88,18 @@ async def resolve_retry_strategy( lock = self._locks[retry_id] async with lock: - return self._create_retry_strategy_cached(retry_id, retry_strategy) + return self._create_retry_strategy_cached( + retry_id, options.retry_mode, options.max_attempts + ) def _create_retry_strategy( - self, retry_strategy: RetryStrategyType + self, retry_mode: RetryStrategyType, max_attempts: int ) -> RetryStrategy: - match retry_strategy: + match retry_mode: case "simple": - return SimpleRetryStrategy() + return SimpleRetryStrategy(max_attempts=max_attempts) case _: - raise ValueError(f"Unknown retry strategy: {retry_strategy}") + raise ValueError(f"Unknown retry mode: {retry_mode}") class ExponentialBackoffJitterType(Enum): diff --git a/packages/smithy-core/tests/unit/test_retries.py b/packages/smithy-core/tests/unit/test_retries.py index 963570ce9..e1dc3cd5a 100644 --- a/packages/smithy-core/tests/unit/test_retries.py +++ b/packages/smithy-core/tests/unit/test_retries.py @@ -1,6 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 import asyncio +from dataclasses import dataclass import pytest from smithy_core.aio.client import CLIENT_ID @@ -8,6 +9,7 @@ from smithy_core.retries import ( CachingRetryStrategyResolver, ExponentialRetryBackoffStrategy, + RetryStrategyOptions, SimpleRetryStrategy, ) from smithy_core.retries import ( @@ -16,6 +18,13 @@ from smithy_core.types import TypedProperties +@dataclass +class MockConfig: + """Mock config for testing retry resolver.""" + + retry_options: RetryStrategyOptions = RetryStrategyOptions() + + @pytest.mark.parametrize( "jitter_type, scale_value, max_backoff, expected_delays", [ @@ -114,7 +123,7 @@ def test_simple_retry_does_not_retry_unsafe() -> None: @pytest.mark.asyncio async def test_caching_retry_strategy_default_resolution() -> None: resolver = CachingRetryStrategyResolver() - properties = TypedProperties({CLIENT_ID.key: "test-client-1"}) + properties = TypedProperties({CLIENT_ID.key: "test-client-1", "config": MockConfig()}) strategy = await resolver.resolve_retry_strategy(properties=properties) @@ -124,8 +133,9 @@ async def test_caching_retry_strategy_default_resolution() -> None: @pytest.mark.asyncio async def test_caching_retry_strategy_resolver_caches_per_client() -> None: resolver = CachingRetryStrategyResolver() - properties1 = TypedProperties({CLIENT_ID.key: "test-caller-1"}) - properties2 = TypedProperties({CLIENT_ID.key: "test-caller-2"}) + config = MockConfig() + properties1 = TypedProperties({CLIENT_ID.key: "test-caller-1", "config": config}) + properties2 = TypedProperties({CLIENT_ID.key: "test-caller-2", "config": config}) strategy1a = await resolver.resolve_retry_strategy(properties=properties1) strategy1b = await resolver.resolve_retry_strategy(properties=properties1) @@ -138,7 +148,7 @@ async def test_caching_retry_strategy_resolver_caches_per_client() -> None: @pytest.mark.asyncio async def test_caching_retry_strategy_resolver_concurrent_access() -> None: resolver = CachingRetryStrategyResolver() - properties = TypedProperties({CLIENT_ID.key: "test-caller-concurrent"}) + properties = TypedProperties({CLIENT_ID.key: "test-caller-concurrent", "config": MockConfig()}) strategies = await asyncio.gather( resolver.resolve_retry_strategy(properties=properties), @@ -150,6 +160,24 @@ async def test_caching_retry_strategy_resolver_concurrent_access() -> None: assert strategies[1] is strategies[2] +@pytest.mark.asyncio +async def test_caching_retry_strategy_resolver_caches_by_options() -> None: + resolver = CachingRetryStrategyResolver() + + config1 = MockConfig(retry_options=RetryStrategyOptions(max_attempts=3)) + config2 = MockConfig(retry_options=RetryStrategyOptions(max_attempts=5)) + + properties1 = TypedProperties({CLIENT_ID.key: "test-client", "config": config1}) + properties2 = TypedProperties({CLIENT_ID.key: "test-client", "config": config2}) + + strategy1 = await resolver.resolve_retry_strategy(properties=properties1) + strategy2 = await resolver.resolve_retry_strategy(properties=properties2) + + assert strategy1 is not strategy2 + assert strategy1.max_attempts == 3 + assert strategy2.max_attempts == 5 + + @pytest.mark.asyncio async def test_caching_retry_strategy_resolver_requires_client_id() -> None: resolver = CachingRetryStrategyResolver() From 829e742c6b743d10527e7505b6dd249ec4648277 Mon Sep 17 00:00:00 2001 From: SamRemis Date: Fri, 7 Nov 2025 15:20:03 -0500 Subject: [PATCH 04/11] run linters --- .../smithy-core/src/smithy_core/retries.py | 6 ++---- .../smithy-core/tests/unit/test_retries.py | 20 +++++++++++-------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index a846df635..170a9e768 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -17,8 +17,7 @@ @dataclass(kw_only=True, frozen=True) class RetryStrategyOptions: - """Options for configuring retry behavior. - """ + """Options for configuring retry behavior.""" retry_mode: RetryStrategyType = "simple" """The retry mode to use.""" @@ -51,8 +50,7 @@ def __init__(self) -> None: self._main_lock = asyncio.Lock() def __deepcopy__(self, memo: dict[int, Any]) -> "CachingRetryStrategyResolver": - """Return self to preserve cache across operation-level config copies. - """ + """Return self to preserve cache across operation-level config copies.""" return self @lru_cache(maxsize=50) diff --git a/packages/smithy-core/tests/unit/test_retries.py b/packages/smithy-core/tests/unit/test_retries.py index e1dc3cd5a..691049ef7 100644 --- a/packages/smithy-core/tests/unit/test_retries.py +++ b/packages/smithy-core/tests/unit/test_retries.py @@ -1,7 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 import asyncio -from dataclasses import dataclass +from dataclasses import dataclass, field import pytest from smithy_core.aio.client import CLIENT_ID @@ -22,7 +22,7 @@ class MockConfig: """Mock config for testing retry resolver.""" - retry_options: RetryStrategyOptions = RetryStrategyOptions() + retry_options: RetryStrategyOptions = field(default_factory=RetryStrategyOptions) @pytest.mark.parametrize( @@ -123,7 +123,9 @@ def test_simple_retry_does_not_retry_unsafe() -> None: @pytest.mark.asyncio async def test_caching_retry_strategy_default_resolution() -> None: resolver = CachingRetryStrategyResolver() - properties = TypedProperties({CLIENT_ID.key: "test-client-1", "config": MockConfig()}) + properties = TypedProperties( + {CLIENT_ID.key: "test-client-1", "config": MockConfig()} + ) strategy = await resolver.resolve_retry_strategy(properties=properties) @@ -148,7 +150,9 @@ async def test_caching_retry_strategy_resolver_caches_per_client() -> None: @pytest.mark.asyncio async def test_caching_retry_strategy_resolver_concurrent_access() -> None: resolver = CachingRetryStrategyResolver() - properties = TypedProperties({CLIENT_ID.key: "test-caller-concurrent", "config": MockConfig()}) + properties = TypedProperties( + {CLIENT_ID.key: "test-caller-concurrent", "config": MockConfig()} + ) strategies = await asyncio.gather( resolver.resolve_retry_strategy(properties=properties), @@ -163,16 +167,16 @@ async def test_caching_retry_strategy_resolver_concurrent_access() -> None: @pytest.mark.asyncio async def test_caching_retry_strategy_resolver_caches_by_options() -> None: resolver = CachingRetryStrategyResolver() - + config1 = MockConfig(retry_options=RetryStrategyOptions(max_attempts=3)) config2 = MockConfig(retry_options=RetryStrategyOptions(max_attempts=5)) - + properties1 = TypedProperties({CLIENT_ID.key: "test-client", "config": config1}) properties2 = TypedProperties({CLIENT_ID.key: "test-client", "config": config2}) - + strategy1 = await resolver.resolve_retry_strategy(properties=properties1) strategy2 = await resolver.resolve_retry_strategy(properties=properties2) - + assert strategy1 is not strategy2 assert strategy1.max_attempts == 3 assert strategy2.max_attempts == 5 From 73b527bbd289240230fa9a1d8f607fbf6a0070ff Mon Sep 17 00:00:00 2001 From: SamRemis Date: Fri, 7 Nov 2025 17:47:33 -0500 Subject: [PATCH 05/11] build-java fixes --- .../amazon/smithy/python/codegen/ClientGenerator.java | 2 +- .../smithy/python/codegen/HttpProtocolTestGenerator.java | 4 ++-- .../smithy/python/codegen/generators/ConfigGenerator.java | 5 +++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index 3090e760e..fd7a78b15 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -208,7 +208,7 @@ raise ExpectationNotMetError("protocol and transport MUST be set on the config t call = ClientCall( input=input, operation=${operation:T}, - context=TypedProperties({"config": config, CLIENT_ID: str(id(self))}), + context=TypedProperties({"config": config, CLIENT_ID.key: str(id(self))}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java index 83f1258f4..00a7b7108 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java @@ -181,13 +181,13 @@ private void generateRequestTest(OperationShape operation, HttpRequestTestCase t } else { path = ""; } - writer.addImport("smithy_core.retries", "SimpleRetryStrategy"); + writer.addImport("smithy_core.retries", "RetryStrategyOptions"); writeClientBlock(context.symbolProvider().toSymbol(service), testCase, Optional.of(() -> { writer.write(""" config = $T( endpoint_uri="https://$L/$L", transport = $T(), - retry_strategy=SimpleRetryStrategy(max_attempts=1), + retry_options=RetryStrategyOptions(max_attempts=1), ) """, CodegenUtils.getConfigSymbol(context.settings()), diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java index dd3842601..e6d14bf65 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java @@ -72,7 +72,8 @@ public final class ConfigGenerator implements Runnable { .initialize(writer -> { writer.addDependency(SmithyPythonDependency.SMITHY_CORE); writer.addImport("smithy_core.retries", "CachingRetryStrategyResolver"); - writer.write("self.retry_strategy_resolver = retry_strategy_resolver or CachingRetryStrategyResolver()"); + writer.write( + "self.retry_strategy_resolver = retry_strategy_resolver or CachingRetryStrategyResolver()"); }) .build(), ConfigProperty.builder() @@ -405,7 +406,7 @@ private void writeInitParams(PythonWriter writer, Collection pro } private void documentProperties(PythonWriter writer, Collection properties) { - writer.writeDocs(() ->{ + writer.writeDocs(() -> { var iter = properties.iterator(); writer.write("\nConstructor.\n"); while (iter.hasNext()) { From 8cd4c921b5282480f1d6a7b1697a6ead5361220c Mon Sep 17 00:00:00 2001 From: SamRemis Date: Thu, 13 Nov 2025 08:36:18 -0500 Subject: [PATCH 06/11] Updates based on feedback --- .../python/codegen/ClientGenerator.java | 29 ++++++- .../codegen/generators/ConfigGenerator.java | 32 +------ .../smithy-core/src/smithy_core/aio/client.py | 14 +-- .../src/smithy_core/interfaces/retries.py | 16 ++-- .../smithy-core/src/smithy_core/retries.py | 67 ++------------- .../smithy-core/tests/unit/test_retries.py | 85 ++----------------- 6 files changed, 59 insertions(+), 184 deletions(-) diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index fd7a78b15..e1802372b 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -83,6 +83,8 @@ private void generateService(PythonWriter writer) { } } + writer.addDependency(SmithyPythonDependency.SMITHY_CORE); + writer.addImport("smithy_core.retries", "CachingRetryStrategyResolver"); writer.write(""" def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): self._config = config or $1T() @@ -95,6 +97,8 @@ def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): for plugin in client_plugins: plugin(self._config) + + self._retry_strategy_resolver = CachingRetryStrategyResolver() """, configSymbol, pluginSymbol, writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins))); var topDownIndex = TopDownIndex.of(model); @@ -183,11 +187,12 @@ private void writeSharedOperationInit(PythonWriter writer, OperationShape operat writer.putContext("operation", symbolProvider.toSymbol(operation)); writer.addImport("smithy_core.aio.client", "ClientCall"); - writer.addImport("smithy_core.aio.client", "CLIENT_ID"); writer.addImport("smithy_core.interceptors", "InterceptorChain"); writer.addImport("smithy_core.types", "TypedProperties"); writer.addImport("smithy_core.aio.client", "RequestPipeline"); writer.addImport("smithy_core.exceptions", "ExpectationNotMetError"); + writer.addImport("smithy_core.retries", "RetryStrategyOptions"); + writer.addImport("smithy_core.interfaces.retries", "RetryStrategy"); writer.addStdlibImport("copy", "deepcopy"); writer.write(""" @@ -201,6 +206,24 @@ private void writeSharedOperationInit(PythonWriter writer, OperationShape operat plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError("protocol and transport MUST be set on the config to make calls.") + + # Resolve retry strategy from config + if isinstance(config.retry_strategy, RetryStrategy): + retry_strategy = config.retry_strategy + elif isinstance(config.retry_strategy, RetryStrategyOptions): + retry_strategy = await self._retry_strategy_resolver.resolve_retry_strategy( + options=config.retry_strategy + ) + elif config.retry_strategy is None: + retry_strategy = await self._retry_strategy_resolver.resolve_retry_strategy( + options=RetryStrategyOptions() + ) + else: + raise TypeError( + f"retry_strategy must be RetryStrategy, RetryStrategyOptions, or None, " + f"got {type(config.retry_strategy).__name__}" + ) + pipeline = RequestPipeline( protocol=config.protocol, transport=config.transport @@ -208,12 +231,12 @@ raise ExpectationNotMetError("protocol and transport MUST be set on the config t call = ClientCall( input=input, operation=${operation:T}, - context=TypedProperties({"config": config, CLIENT_ID.key: str(id(self))}), + context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, - retry_strategy_resolver=config.retry_strategy_resolver, + retry_strategy=retry_strategy, ) """, writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins))); diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java index e6d14bf65..2b788e47f 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java @@ -53,46 +53,22 @@ public final class ConfigGenerator implements Runnable { .initialize(writer -> writer.write("self.interceptors = interceptors or []")) .build(), ConfigProperty.builder() - .name("retry_strategy_resolver") + .name("retry_strategy") .type(Symbol.builder() - .name("RetryStrategyResolver[RetryStrategy]") - .addReference(Symbol.builder() - .name("RetryStrategyResolver") - .namespace("smithy_core.interfaces.retries", ".") - .addDependency(SmithyPythonDependency.SMITHY_CORE) - .build()) + .name("RetryStrategy | RetryStrategyOptions") .addReference(Symbol.builder() .name("RetryStrategy") .namespace("smithy_core.interfaces.retries", ".") .addDependency(SmithyPythonDependency.SMITHY_CORE) .build()) - .build()) - .documentation("The retry strategy resolver for resolving retry strategies per client.") - .nullable(false) - .initialize(writer -> { - writer.addDependency(SmithyPythonDependency.SMITHY_CORE); - writer.addImport("smithy_core.retries", "CachingRetryStrategyResolver"); - writer.write( - "self.retry_strategy_resolver = retry_strategy_resolver or CachingRetryStrategyResolver()"); - }) - .build(), - ConfigProperty.builder() - .name("retry_options") - .type(Symbol.builder() - .name("RetryStrategyOptions") .addReference(Symbol.builder() .name("RetryStrategyOptions") .namespace("smithy_core.retries", ".") .addDependency(SmithyPythonDependency.SMITHY_CORE) .build()) .build()) - .documentation("Options for configuring retry behavior.") - .nullable(false) - .initialize(writer -> { - writer.addDependency(SmithyPythonDependency.SMITHY_CORE); - writer.addImport("smithy_core.retries", "RetryStrategyOptions"); - writer.write("self.retry_options = retry_options or RetryStrategyOptions()"); - }) + .documentation( + "The retry strategy or options for configuring retry behavior. Can be either a fully configured RetryStrategy or RetryStrategyOptions to create one.") .build(), ConfigProperty.builder() .name("endpoint_uri") diff --git a/packages/smithy-core/src/smithy_core/aio/client.py b/packages/smithy-core/src/smithy_core/aio/client.py index 08a061588..bf27c440c 100644 --- a/packages/smithy-core/src/smithy_core/aio/client.py +++ b/packages/smithy-core/src/smithy_core/aio/client.py @@ -22,7 +22,7 @@ ) from ..interfaces import Endpoint, TypedProperties from ..interfaces.auth import AuthOption, AuthSchemeResolver -from ..interfaces.retries import RetryStrategy, RetryStrategyResolver +from ..interfaces.retries import RetryStrategy from ..schemas import APIOperation from ..serializers import SerializeableShape from ..shapes import ShapeID @@ -44,10 +44,6 @@ AUTH_SCHEME = PropertyKey(key="auth_scheme", value_type=AuthScheme[Any, Any, Any, Any]) -CLIENT_ID = PropertyKey(key="client_id", value_type=str) -"""A unique identifier for the client instance. -""" - _UNRESOLVED = URI(host="", path="/") _LOGGER = logging.getLogger(__name__) @@ -81,8 +77,8 @@ class ClientCall[I: SerializeableShape, O: DeserializeableShape]: endpoint_resolver: EndpointResolver """The endpoint resolver for the operation.""" - retry_strategy_resolver: RetryStrategyResolver[RetryStrategy] - """The retry strategy resolver for the operation.""" + retry_strategy: RetryStrategy + """The retry strategy to use for the operation.""" retry_scope: str | None = None """The retry scope for the operation.""" @@ -333,9 +329,7 @@ async def _retry[I: SerializeableShape, O: DeserializeableShape]( if not call.retryable(): return await self._handle_attempt(call, request_context, request_future) - retry_strategy = await call.retry_strategy_resolver.resolve_retry_strategy( - properties=request_context.properties - ) + retry_strategy = call.retry_strategy retry_token = retry_strategy.acquire_initial_retry_token( token_scope=call.retry_scope ) diff --git a/packages/smithy-core/src/smithy_core/interfaces/retries.py b/packages/smithy-core/src/smithy_core/interfaces/retries.py index 374d10fd2..9f7438bdd 100644 --- a/packages/smithy-core/src/smithy_core/interfaces/retries.py +++ b/packages/smithy-core/src/smithy_core/interfaces/retries.py @@ -1,9 +1,12 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + from dataclasses import dataclass -from typing import Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable -from . import TypedProperties +if TYPE_CHECKING: + from smithy_core.retries import RetryStrategyOptions @runtime_checkable @@ -57,6 +60,7 @@ class RetryToken(Protocol): RetryStrategyType = Literal["simple"] +@runtime_checkable class RetryStrategy(Protocol): """Issuer of :py:class:`RetryToken`s.""" @@ -108,11 +112,11 @@ def record_success(self, *, token: RetryToken) -> None: class RetryStrategyResolver[RS: RetryStrategy](Protocol): - """Used to resolve a RetryStrategy for a given caller.""" + """Used to resolve a RetryStrategy from retry options.""" - async def resolve_retry_strategy(self, *, properties: TypedProperties) -> RS: - """Resolve the retry strategy for the caller. + async def resolve_retry_strategy(self, *, options: RetryStrategyOptions) -> RS: + """Resolve the retry strategy from the provided options. - :param properties: Properties including caller identification and config. + :param options: The retry strategy options to use for creating the strategy. """ ... diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index 170a9e768..b25a0425c 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -1,18 +1,13 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import asyncio import random from collections.abc import Callable from dataclasses import dataclass from enum import Enum -from functools import lru_cache -from typing import Any, Protocol -from .aio.client import CLIENT_ID from .exceptions import RetryError from .interfaces import retries as retries_interface from .interfaces.retries import RetryStrategy, RetryStrategyResolver, RetryStrategyType -from .types import PropertyKey @dataclass(kw_only=True, frozen=True) @@ -26,69 +21,21 @@ class RetryStrategyOptions: """Maximum number of attempts (initial attempt plus retries).""" -class RetryConfig(Protocol): - """Protocol for config objects that support retry configuration.""" - - retry_options: RetryStrategyOptions - - -RETRY_CONFIG = PropertyKey(key="config", value_type=RetryConfig) - - class CachingRetryStrategyResolver(RetryStrategyResolver[RetryStrategy]): - """Caching retry strategy resolver that creates and caches retry strategies per caller. - - This resolver maintains a cache of retry strategies keyed by a unique identifier - for each caller. This allows multiple operations from the same caller to share - a single retry strategy instance, which is important for strategies that maintain - state across retries (e.g., token buckets, rate limiters). + """Retry strategy resolver that creates retry strategies based on configuration options. + This resolver creates a new retry strategy instance for each resolution request + based on the retry options provided. """ - def __init__(self) -> None: - self._locks: dict[str, asyncio.Lock] = {} - self._main_lock = asyncio.Lock() - - def __deepcopy__(self, memo: dict[int, Any]) -> "CachingRetryStrategyResolver": - """Return self to preserve cache across operation-level config copies.""" - return self - - @lru_cache(maxsize=50) - def _create_retry_strategy_cached( - self, retry_id: str, retry_mode: RetryStrategyType, max_attempts: int - ) -> RetryStrategy: - return self._create_retry_strategy(retry_mode, max_attempts) - async def resolve_retry_strategy( - self, *, properties: retries_interface.TypedProperties + self, *, options: RetryStrategyOptions ) -> RetryStrategy: - """Get or create a retry strategy for the caller. + """Create a retry strategy based on the provided options. - :param properties: Properties map that must contain the CLIENT_ID property key - with a unique identifier for the caller, and a "config" key with a - retry_strategy attribute (RetryStrategyOptions) specifying the strategy - configuration. Strategies are cached per client and options combination. - :raises ValueError: If CLIENT_ID is not present in properties. + :param options: The retry strategy options to use for creating the strategy. """ - retry_id = properties.get(CLIENT_ID.key) - if retry_id is None: - raise ValueError( - f"Properties must contain '{CLIENT_ID.key}' key with a unique identifier for the caller" - ) - - # Get retry options from config - config = properties[RETRY_CONFIG] - options = config.retry_options - - async with self._main_lock: - if retry_id not in self._locks: - self._locks[retry_id] = asyncio.Lock() - lock = self._locks[retry_id] - - async with lock: - return self._create_retry_strategy_cached( - retry_id, options.retry_mode, options.max_attempts - ) + return self._create_retry_strategy(options.retry_mode, options.max_attempts) def _create_retry_strategy( self, retry_mode: RetryStrategyType, max_attempts: int diff --git a/packages/smithy-core/tests/unit/test_retries.py b/packages/smithy-core/tests/unit/test_retries.py index 691049ef7..7176f9622 100644 --- a/packages/smithy-core/tests/unit/test_retries.py +++ b/packages/smithy-core/tests/unit/test_retries.py @@ -1,10 +1,6 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import asyncio -from dataclasses import dataclass, field - import pytest -from smithy_core.aio.client import CLIENT_ID from smithy_core.exceptions import CallError, RetryError from smithy_core.retries import ( CachingRetryStrategyResolver, @@ -15,14 +11,6 @@ from smithy_core.retries import ( ExponentialBackoffJitterType as EBJT, ) -from smithy_core.types import TypedProperties - - -@dataclass -class MockConfig: - """Mock config for testing retry resolver.""" - - retry_options: RetryStrategyOptions = field(default_factory=RetryStrategyOptions) @pytest.mark.parametrize( @@ -120,81 +108,24 @@ def test_simple_retry_does_not_retry_unsafe() -> None: strategy.refresh_retry_token_for_retry(token_to_renew=token, error=error) -@pytest.mark.asyncio async def test_caching_retry_strategy_default_resolution() -> None: resolver = CachingRetryStrategyResolver() - properties = TypedProperties( - {CLIENT_ID.key: "test-client-1", "config": MockConfig()} - ) + options = RetryStrategyOptions() - strategy = await resolver.resolve_retry_strategy(properties=properties) + strategy = await resolver.resolve_retry_strategy(options=options) assert isinstance(strategy, SimpleRetryStrategy) + assert strategy.max_attempts == 3 -@pytest.mark.asyncio -async def test_caching_retry_strategy_resolver_caches_per_client() -> None: - resolver = CachingRetryStrategyResolver() - config = MockConfig() - properties1 = TypedProperties({CLIENT_ID.key: "test-caller-1", "config": config}) - properties2 = TypedProperties({CLIENT_ID.key: "test-caller-2", "config": config}) - - strategy1a = await resolver.resolve_retry_strategy(properties=properties1) - strategy1b = await resolver.resolve_retry_strategy(properties=properties1) - strategy2 = await resolver.resolve_retry_strategy(properties=properties2) - - assert strategy1a is strategy1b - assert strategy1a is not strategy2 - - -@pytest.mark.asyncio -async def test_caching_retry_strategy_resolver_concurrent_access() -> None: +async def test_caching_retry_strategy_resolver_creates_strategies_by_options() -> None: resolver = CachingRetryStrategyResolver() - properties = TypedProperties( - {CLIENT_ID.key: "test-caller-concurrent", "config": MockConfig()} - ) - strategies = await asyncio.gather( - resolver.resolve_retry_strategy(properties=properties), - resolver.resolve_retry_strategy(properties=properties), - resolver.resolve_retry_strategy(properties=properties), - ) - - assert strategies[0] is strategies[1] - assert strategies[1] is strategies[2] + options1 = RetryStrategyOptions(max_attempts=3) + options2 = RetryStrategyOptions(max_attempts=5) + strategy1 = await resolver.resolve_retry_strategy(options=options1) + strategy2 = await resolver.resolve_retry_strategy(options=options2) -@pytest.mark.asyncio -async def test_caching_retry_strategy_resolver_caches_by_options() -> None: - resolver = CachingRetryStrategyResolver() - - config1 = MockConfig(retry_options=RetryStrategyOptions(max_attempts=3)) - config2 = MockConfig(retry_options=RetryStrategyOptions(max_attempts=5)) - - properties1 = TypedProperties({CLIENT_ID.key: "test-client", "config": config1}) - properties2 = TypedProperties({CLIENT_ID.key: "test-client", "config": config2}) - - strategy1 = await resolver.resolve_retry_strategy(properties=properties1) - strategy2 = await resolver.resolve_retry_strategy(properties=properties2) - - assert strategy1 is not strategy2 assert strategy1.max_attempts == 3 assert strategy2.max_attempts == 5 - - -@pytest.mark.asyncio -async def test_caching_retry_strategy_resolver_requires_client_id() -> None: - resolver = CachingRetryStrategyResolver() - properties = TypedProperties({}) - - with pytest.raises(ValueError, match=CLIENT_ID.key): - await resolver.resolve_retry_strategy(properties=properties) - - -def test_caching_retry_strategy_resolver_survives_deepcopy() -> None: - from copy import deepcopy - - resolver = CachingRetryStrategyResolver() - resolver_copy = deepcopy(resolver) - - assert resolver is resolver_copy From 78700548386c4e4c15f156fea95515aade94210d Mon Sep 17 00:00:00 2001 From: SamRemis Date: Thu, 13 Nov 2025 13:48:14 -0500 Subject: [PATCH 07/11] Update HttpProtocolTestGenerator.java --- .../amazon/smithy/python/codegen/HttpProtocolTestGenerator.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java index 00a7b7108..54a0d6028 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java @@ -187,7 +187,7 @@ private void generateRequestTest(OperationShape operation, HttpRequestTestCase t config = $T( endpoint_uri="https://$L/$L", transport = $T(), - retry_options=RetryStrategyOptions(max_attempts=1), + retry_strategy=RetryStrategyOptions(max_attempts=1), ) """, CodegenUtils.getConfigSymbol(context.settings()), From ff2c36814a5515f23483cf09318710d9fb308f7e Mon Sep 17 00:00:00 2001 From: SamRemis Date: Thu, 13 Nov 2025 14:18:25 -0500 Subject: [PATCH 08/11] Add back cache and lock --- .../smithy-core/src/smithy_core/retries.py | 20 ++++++++++++++----- .../smithy-core/tests/unit/test_retries.py | 10 ++++++++++ 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index b25a0425c..baf75fe76 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -1,6 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 import random +import threading from collections.abc import Callable from dataclasses import dataclass from enum import Enum @@ -22,20 +23,29 @@ class RetryStrategyOptions: class CachingRetryStrategyResolver(RetryStrategyResolver[RetryStrategy]): - """Retry strategy resolver that creates retry strategies based on configuration options. + """Retry strategy resolver that caches retry strategies based on configuration options. - This resolver creates a new retry strategy instance for each resolution request - based on the retry options provided. + This resolver caches retry strategy instances based on their configuration to reuse existing + instances of RetryStrategy with the same settings. """ + def __init__(self) -> None: + self._cache: dict[RetryStrategyOptions, RetryStrategy] = {} + self._lock = threading.Lock() + async def resolve_retry_strategy( self, *, options: RetryStrategyOptions ) -> RetryStrategy: - """Create a retry strategy based on the provided options. + """Resolve a retry strategy from the provided options, using cache when possible. :param options: The retry strategy options to use for creating the strategy. """ - return self._create_retry_strategy(options.retry_mode, options.max_attempts) + with self._lock: + if options not in self._cache: + self._cache[options] = self._create_retry_strategy( + options.retry_mode, options.max_attempts + ) + return self._cache[options] def _create_retry_strategy( self, retry_mode: RetryStrategyType, max_attempts: int diff --git a/packages/smithy-core/tests/unit/test_retries.py b/packages/smithy-core/tests/unit/test_retries.py index 7176f9622..b6460c4a3 100644 --- a/packages/smithy-core/tests/unit/test_retries.py +++ b/packages/smithy-core/tests/unit/test_retries.py @@ -129,3 +129,13 @@ async def test_caching_retry_strategy_resolver_creates_strategies_by_options() - assert strategy1.max_attempts == 3 assert strategy2.max_attempts == 5 + + +async def test_caching_retry_strategy_resolver_caches_strategies() -> None: + resolver = CachingRetryStrategyResolver() + + options = RetryStrategyOptions(max_attempts=5) + strategy1 = await resolver.resolve_retry_strategy(options=options) + strategy2 = await resolver.resolve_retry_strategy(options=options) + + assert strategy1 is strategy2 From dd748a89e7ec972d2075e2cdc964592745929b34 Mon Sep 17 00:00:00 2001 From: SamRemis Date: Thu, 13 Nov 2025 19:19:00 -0500 Subject: [PATCH 09/11] More updates with feedback --- .../python/codegen/ClientGenerator.java | 4 +-- .../codegen/HttpProtocolTestGenerator.java | 4 +-- .../src/smithy_core/interfaces/retries.py | 21 +-------------- .../smithy-core/src/smithy_core/retries.py | 26 ++++++++----------- .../smithy-core/tests/unit/test_retries.py | 8 +++--- 5 files changed, 20 insertions(+), 43 deletions(-) diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index e1802372b..95c6e0d5a 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -84,7 +84,7 @@ private void generateService(PythonWriter writer) { } writer.addDependency(SmithyPythonDependency.SMITHY_CORE); - writer.addImport("smithy_core.retries", "CachingRetryStrategyResolver"); + writer.addImport("smithy_core.retries", "RetryStrategyResolver"); writer.write(""" def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): self._config = config or $1T() @@ -98,7 +98,7 @@ def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): for plugin in client_plugins: plugin(self._config) - self._retry_strategy_resolver = CachingRetryStrategyResolver() + self._retry_strategy_resolver = RetryStrategyResolver() """, configSymbol, pluginSymbol, writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins))); var topDownIndex = TopDownIndex.of(model); diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java index 54a0d6028..83f1258f4 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java @@ -181,13 +181,13 @@ private void generateRequestTest(OperationShape operation, HttpRequestTestCase t } else { path = ""; } - writer.addImport("smithy_core.retries", "RetryStrategyOptions"); + writer.addImport("smithy_core.retries", "SimpleRetryStrategy"); writeClientBlock(context.symbolProvider().toSymbol(service), testCase, Optional.of(() -> { writer.write(""" config = $T( endpoint_uri="https://$L/$L", transport = $T(), - retry_strategy=RetryStrategyOptions(max_attempts=1), + retry_strategy=SimpleRetryStrategy(max_attempts=1), ) """, CodegenUtils.getConfigSymbol(context.settings()), diff --git a/packages/smithy-core/src/smithy_core/interfaces/retries.py b/packages/smithy-core/src/smithy_core/interfaces/retries.py index 9f7438bdd..19e96362a 100644 --- a/packages/smithy-core/src/smithy_core/interfaces/retries.py +++ b/packages/smithy-core/src/smithy_core/interfaces/retries.py @@ -1,12 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - from dataclasses import dataclass -from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable - -if TYPE_CHECKING: - from smithy_core.retries import RetryStrategyOptions +from typing import Protocol, runtime_checkable @runtime_checkable @@ -57,9 +52,6 @@ class RetryToken(Protocol): """Delay in seconds to wait before the retry attempt.""" -RetryStrategyType = Literal["simple"] - - @runtime_checkable class RetryStrategy(Protocol): """Issuer of :py:class:`RetryToken`s.""" @@ -109,14 +101,3 @@ def record_success(self, *, token: RetryToken) -> None: :param token: The token used for the previous successful attempt. """ ... - - -class RetryStrategyResolver[RS: RetryStrategy](Protocol): - """Used to resolve a RetryStrategy from retry options.""" - - async def resolve_retry_strategy(self, *, options: RetryStrategyOptions) -> RS: - """Resolve the retry strategy from the provided options. - - :param options: The retry strategy options to use for creating the strategy. - """ - ... diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index baf75fe76..e882fa782 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -1,14 +1,17 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 import random -import threading from collections.abc import Callable from dataclasses import dataclass from enum import Enum +from functools import lru_cache +from typing import Literal from .exceptions import RetryError from .interfaces import retries as retries_interface -from .interfaces.retries import RetryStrategy, RetryStrategyResolver, RetryStrategyType +from .interfaces.retries import RetryStrategy + +RetryStrategyType = Literal["simple"] @dataclass(kw_only=True, frozen=True) @@ -22,17 +25,13 @@ class RetryStrategyOptions: """Maximum number of attempts (initial attempt plus retries).""" -class CachingRetryStrategyResolver(RetryStrategyResolver[RetryStrategy]): +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. + instances of RetryStrategy with the same settings. Uses LRU cache for thread-safe caching. """ - def __init__(self) -> None: - self._cache: dict[RetryStrategyOptions, RetryStrategy] = {} - self._lock = threading.Lock() - async def resolve_retry_strategy( self, *, options: RetryStrategyOptions ) -> RetryStrategy: @@ -40,15 +39,12 @@ async def resolve_retry_strategy( :param options: The retry strategy options to use for creating the strategy. """ - with self._lock: - if options not in self._cache: - self._cache[options] = self._create_retry_strategy( - options.retry_mode, options.max_attempts - ) - return self._cache[options] + return self._create_retry_strategy(options.retry_mode, options.max_attempts) + @staticmethod + @lru_cache def _create_retry_strategy( - self, retry_mode: RetryStrategyType, max_attempts: int + retry_mode: RetryStrategyType, max_attempts: int ) -> RetryStrategy: match retry_mode: case "simple": diff --git a/packages/smithy-core/tests/unit/test_retries.py b/packages/smithy-core/tests/unit/test_retries.py index b6460c4a3..43d8db0d9 100644 --- a/packages/smithy-core/tests/unit/test_retries.py +++ b/packages/smithy-core/tests/unit/test_retries.py @@ -3,7 +3,7 @@ import pytest from smithy_core.exceptions import CallError, RetryError from smithy_core.retries import ( - CachingRetryStrategyResolver, + RetryStrategyResolver, ExponentialRetryBackoffStrategy, RetryStrategyOptions, SimpleRetryStrategy, @@ -109,7 +109,7 @@ def test_simple_retry_does_not_retry_unsafe() -> None: async def test_caching_retry_strategy_default_resolution() -> None: - resolver = CachingRetryStrategyResolver() + resolver = RetryStrategyResolver() options = RetryStrategyOptions() strategy = await resolver.resolve_retry_strategy(options=options) @@ -119,7 +119,7 @@ async def test_caching_retry_strategy_default_resolution() -> None: async def test_caching_retry_strategy_resolver_creates_strategies_by_options() -> None: - resolver = CachingRetryStrategyResolver() + resolver = RetryStrategyResolver() options1 = RetryStrategyOptions(max_attempts=3) options2 = RetryStrategyOptions(max_attempts=5) @@ -132,7 +132,7 @@ async def test_caching_retry_strategy_resolver_creates_strategies_by_options() - async def test_caching_retry_strategy_resolver_caches_strategies() -> None: - resolver = CachingRetryStrategyResolver() + resolver = RetryStrategyResolver() options = RetryStrategyOptions(max_attempts=5) strategy1 = await resolver.resolve_retry_strategy(options=options) From 572302eff1632830cf053b91b4d82b26e91674ef Mon Sep 17 00:00:00 2001 From: SamRemis Date: Thu, 13 Nov 2025 19:37:03 -0500 Subject: [PATCH 10/11] Remove staticmethod --- packages/smithy-core/src/smithy_core/retries.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index e882fa782..78d139d90 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -41,10 +41,9 @@ async def resolve_retry_strategy( """ return self._create_retry_strategy(options.retry_mode, options.max_attempts) - @staticmethod @lru_cache def _create_retry_strategy( - retry_mode: RetryStrategyType, max_attempts: int + self, retry_mode: RetryStrategyType, max_attempts: int ) -> RetryStrategy: match retry_mode: case "simple": From 24aa35b851055e71f8309ec655d4b68e7458d9cd Mon Sep 17 00:00:00 2001 From: SamRemis Date: Thu, 13 Nov 2025 19:40:10 -0500 Subject: [PATCH 11/11] run linter --- packages/smithy-core/tests/unit/test_retries.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/smithy-core/tests/unit/test_retries.py b/packages/smithy-core/tests/unit/test_retries.py index 43d8db0d9..fd87a789b 100644 --- a/packages/smithy-core/tests/unit/test_retries.py +++ b/packages/smithy-core/tests/unit/test_retries.py @@ -3,14 +3,14 @@ import pytest from smithy_core.exceptions import CallError, RetryError from smithy_core.retries import ( - RetryStrategyResolver, + ExponentialBackoffJitterType as EBJT, +) +from smithy_core.retries import ( ExponentialRetryBackoffStrategy, RetryStrategyOptions, + RetryStrategyResolver, SimpleRetryStrategy, ) -from smithy_core.retries import ( - ExponentialBackoffJitterType as EBJT, -) @pytest.mark.parametrize(