Skip to content

Commit 312d190

Browse files
Zelda Hesslerjdisanti
andauthored
Update standard orchestrator retries with token bucket and more tests (#2764)
## Motivation and Context <!--- Why is this change required? What problem does it solve? --> <!--- If it fixes an open issue, please link to the issue here --> addresses #2743 ## Description <!--- Describe your changes in detail --> - add more standard retry tests - add optional standard retries token bucket ## Testing <!--- Please describe in detail how you tested your changes --> <!--- Include details of your testing environment, and the tests you ran to --> <!--- see how your change affects other areas of the code, etc. --> tests are included ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._ --------- Co-authored-by: John DiSanti <jdisanti@amazon.com>
1 parent 5473192 commit 312d190

File tree

16 files changed

+655
-455
lines changed

16 files changed

+655
-455
lines changed

aws/rust-runtime/aws-config/src/profile/credentials.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -465,7 +465,6 @@ async fn build_provider_chain(
465465

466466
#[cfg(test)]
467467
mod test {
468-
469468
use crate::profile::credentials::Builder;
470469
use crate::test_case::TestEnvironment;
471470

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationGenerator.kt

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -146,20 +146,6 @@ open class OperationGenerator(
146146
if (codegenContext.smithyRuntimeMode.generateOrchestrator) {
147147
rustTemplate(
148148
"""
149-
pub(crate) fn register_runtime_plugins(
150-
runtime_plugins: #{RuntimePlugins},
151-
handle: #{Arc}<crate::client::Handle>,
152-
config_override: #{Option}<crate::config::Builder>,
153-
) -> #{RuntimePlugins} {
154-
#{register_default_runtime_plugins}(
155-
runtime_plugins,
156-
#{Box}::new(Self::new()) as _,
157-
handle,
158-
config_override
159-
)
160-
#{additional_runtime_plugins}
161-
}
162-
163149
pub(crate) async fn orchestrate(
164150
runtime_plugins: &#{RuntimePlugins},
165151
input: #{Input},
@@ -186,6 +172,20 @@ open class OperationGenerator(
186172
let input = #{TypedBox}::new(input).erase();
187173
#{invoke_with_stop_point}(input, runtime_plugins, stop_point).await
188174
}
175+
176+
pub(crate) fn register_runtime_plugins(
177+
runtime_plugins: #{RuntimePlugins},
178+
handle: #{Arc}<crate::client::Handle>,
179+
config_override: #{Option}<crate::config::Builder>,
180+
) -> #{RuntimePlugins} {
181+
#{register_default_runtime_plugins}(
182+
runtime_plugins,
183+
#{Box}::new(Self::new()) as _,
184+
handle,
185+
config_override
186+
)
187+
#{additional_runtime_plugins}
188+
}
189189
""",
190190
*codegenScope,
191191
"Error" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::interceptors::context::Error"),

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ServiceRuntimePluginGenerator.kt

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,8 @@ class ServiceRuntimePluginGenerator(
105105
"StaticAuthOptionResolver" to runtimeApi.resolve("client::auth::option_resolver::StaticAuthOptionResolver"),
106106
"default_connector" to client.resolve("conns::default_connector"),
107107
"require_connector" to client.resolve("conns::require_connector"),
108+
"TimeoutConfig" to smithyTypes.resolve("timeout::TimeoutConfig"),
109+
"RetryConfig" to smithyTypes.resolve("retry::RetryConfig"),
108110
)
109111
}
110112

@@ -142,20 +144,17 @@ class ServiceRuntimePluginGenerator(
142144
self.handle.conf.endpoint_resolver());
143145
cfg.set_endpoint_resolver(endpoint_resolver);
144146
145-
// TODO(enableNewSmithyRuntime): Use the `store_append` method of ConfigBag to insert classifiers
146-
let retry_classifiers = #{RetryClassifiers}::new()
147-
#{retry_classifier_customizations};
148-
cfg.set_retry_classifiers(retry_classifiers);
147+
// TODO(enableNewSmithyRuntime): Make it possible to set retry classifiers at the service level.
148+
// Retry classifiers can also be set at the operation level and those should be added to the
149+
// list of classifiers defined here, rather than replacing them.
149150
150151
let sleep_impl = self.handle.conf.sleep_impl();
151-
let timeout_config = self.handle.conf.timeout_config();
152-
let retry_config = self.handle.conf.retry_config();
152+
let timeout_config = self.handle.conf.timeout_config().cloned().unwrap_or_else(|| #{TimeoutConfig}::disabled());
153+
let retry_config = self.handle.conf.retry_config().cloned().unwrap_or_else(|| #{RetryConfig}::disabled());
153154
154-
if let Some(retry_config) = retry_config {
155-
cfg.set_retry_strategy(#{StandardRetryStrategy}::new(retry_config));
156-
}
155+
cfg.set_retry_strategy(#{StandardRetryStrategy}::new(&retry_config));
157156
158-
let connector_settings = timeout_config.map(#{ConnectorSettings}::from_timeout_config).unwrap_or_default();
157+
let connector_settings = #{ConnectorSettings}::from_timeout_config(&timeout_config);
159158
if let Some(connection) = self.handle.conf.http_connector()
160159
.and_then(|c| c.connector(&connector_settings, sleep_impl.clone()))
161160
.or_else(|| #{default_connector}(&connector_settings, sleep_impl)) {
@@ -180,9 +179,6 @@ class ServiceRuntimePluginGenerator(
180179
"http_auth_scheme_customizations" to writable {
181180
writeCustomizations(customizations, ServiceRuntimePluginSection.HttpAuthScheme("cfg"))
182181
},
183-
"retry_classifier_customizations" to writable {
184-
writeCustomizations(customizations, ServiceRuntimePluginSection.RetryClassifier("cfg"))
185-
},
186182
"additional_config" to writable {
187183
writeCustomizations(customizations, ServiceRuntimePluginSection.AdditionalConfig("cfg"))
188184
},

rust-runtime/aws-smithy-runtime-api/src/client/orchestrator/error.rs

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,12 @@ pub enum OrchestratorError<E> {
1818
Interceptor { err: InterceptorError },
1919
/// An error returned by a service.
2020
Operation { err: E },
21+
/// An error that occurs when a request times out.
22+
Timeout { err: BoxError },
23+
/// An error that occurs when request dispatch fails.
24+
Connector { err: ConnectorError },
25+
/// An error that occurs when a response can't be deserialized.
26+
Response { err: BoxError },
2127
/// A general orchestrator error.
2228
Other { err: BoxError },
2329
}
@@ -34,11 +40,26 @@ impl<E: Debug> OrchestratorError<E> {
3440
Self::Operation { err }
3541
}
3642

37-
/// Create a new `OrchestratorError` from an [`InterceptorError`].
43+
/// Create a new `OrchestratorError::Interceptor` from an [`InterceptorError`].
3844
pub fn interceptor(err: InterceptorError) -> Self {
3945
Self::Interceptor { err }
4046
}
4147

48+
/// Create a new `OrchestratorError::Timeout` from a [`BoxError`].
49+
pub fn timeout(err: BoxError) -> Self {
50+
Self::Timeout { err }
51+
}
52+
53+
/// Create a new `OrchestratorError::Response` from a [`BoxError`].
54+
pub fn response(err: BoxError) -> Self {
55+
Self::Response { err }
56+
}
57+
58+
/// Create a new `OrchestratorError::Connector` from a [`ConnectorError`].
59+
pub fn connector(err: ConnectorError) -> Self {
60+
Self::Connector { err }
61+
}
62+
4263
/// Convert the `OrchestratorError` into `Some` operation specific error if it is one. Otherwise,
4364
/// return `None`.
4465
pub fn as_operation_error(&self) -> Option<&E> {
@@ -72,6 +93,9 @@ impl<E: Debug> OrchestratorError<E> {
7293
debug_assert!(phase.is_after_deserialization(), "operation errors are a result of successfully receiving and parsing a response from the server. Therefore, we must be in the 'After Deserialization' phase.");
7394
SdkError::service_error(err, response.expect("phase has a response"))
7495
}
96+
Self::Connector { err } => SdkError::dispatch_failure(err),
97+
Self::Timeout { err } => SdkError::timeout_error(err),
98+
Self::Response { err } => SdkError::response_error(err, response.unwrap()),
7599
Self::Other { err } => {
76100
use Phase::*;
77101
match phase {
@@ -111,15 +135,6 @@ where
111135
}
112136
}
113137

114-
impl<E> From<BoxError> for OrchestratorError<E>
115-
where
116-
E: Debug + std::error::Error + 'static,
117-
{
118-
fn from(err: BoxError) -> Self {
119-
Self::other(err)
120-
}
121-
}
122-
123138
impl From<TypeErasedError> for OrchestratorError<TypeErasedError> {
124139
fn from(err: TypeErasedError) -> Self {
125140
Self::operation(err)

rust-runtime/aws-smithy-runtime-api/src/client/retries.rs

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,16 @@ pub enum ShouldAttempt {
2020
YesAfterDelay(Duration),
2121
}
2222

23+
#[cfg(feature = "test-util")]
24+
impl ShouldAttempt {
25+
pub fn expect_delay(self) -> Duration {
26+
match self {
27+
ShouldAttempt::YesAfterDelay(delay) => delay,
28+
_ => panic!("Expected this to be the `YesAfterDelay` variant but it was the `{self:?}` variant instead"),
29+
}
30+
}
31+
}
32+
2333
pub trait RetryStrategy: Send + Sync + Debug {
2434
fn should_attempt_initial_request(&self, cfg: &ConfigBag) -> Result<ShouldAttempt, BoxError>;
2535

@@ -31,7 +41,7 @@ pub trait RetryStrategy: Send + Sync + Debug {
3141
}
3242

3343
#[non_exhaustive]
34-
#[derive(Eq, PartialEq, Debug)]
44+
#[derive(Clone, Eq, PartialEq, Debug)]
3545
pub enum RetryReason {
3646
Error(ErrorKind),
3747
Explicit(Duration),
@@ -72,10 +82,10 @@ impl RetryClassifiers {
7282
}
7383

7484
impl ClassifyRetry for RetryClassifiers {
75-
fn classify_retry(&self, error: &InterceptorContext) -> Option<RetryReason> {
85+
fn classify_retry(&self, ctx: &InterceptorContext) -> Option<RetryReason> {
7686
// return the first non-None result
7787
self.inner.iter().find_map(|cr| {
78-
let maybe_reason = cr.classify_retry(error);
88+
let maybe_reason = cr.classify_retry(ctx);
7989

8090
match maybe_reason.as_ref() {
8191
Some(reason) => trace!(

rust-runtime/aws-smithy-runtime-api/src/client/retries/rate_limiting.rs

Lines changed: 0 additions & 13 deletions
This file was deleted.

rust-runtime/aws-smithy-runtime-api/src/client/retries/rate_limiting/error.rs

Lines changed: 0 additions & 50 deletions
This file was deleted.

rust-runtime/aws-smithy-runtime-api/src/client/retries/rate_limiting/token.rs

Lines changed: 0 additions & 65 deletions
This file was deleted.

0 commit comments

Comments
 (0)