Skip to content

Commit

Permalink
feature: orchestrator retry classifiers (#2621)
Browse files Browse the repository at this point in the history
## 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 -->
To retry a response, we must first classify it as retryable.

## Description
<!--- Describe your changes in detail -->
feature: add AWS error code classifier
feature: add x-amz-retry-after header classifier
feature: add smithy modeled retry classifier
feature: add error type classifier
feature: add HTTP status code classifier
add: tests for classifiers
remove: redundant `http` dep from `aws-http`
move: `NeverRetryStrategy` to smithy-runtime crate
add: RuntimePluginImpls codegen section for operation-specific runtime
plugin definitions
update: orchestrator retries to work with `ShouldAttempt`
add: retry classifier config bag accessor
add: raw response getter to SdkError
update: RetryStrategy trait signatures to use `ShouldAttempt`
add: `RetryClassifiers` struct for holding and calling retry classifiers
update: `RetryClassifierDecorator` to define orchestrator classifiers
add: `default_retry_classifiers` fn to codegen
update: `ServiceGenerator` to add feature flag for
aws-smithy-runtime/test-util
update: SRA integration test to insert retry classifier plugin

## 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. -->
this change includes tests

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._
  • Loading branch information
Velfi authored Apr 25, 2023
1 parent 21249b0 commit ae995fb
Show file tree
Hide file tree
Showing 24 changed files with 774 additions and 69 deletions.
1 change: 0 additions & 1 deletion aws/rust-runtime/aws-http/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@ aws-smithy-checksums = { path = "../../../rust-runtime/aws-smithy-checksums" }
aws-smithy-protocol-test = { path = "../../../rust-runtime/aws-smithy-protocol-test" }
bytes-utils = "0.1.2"
env_logger = "0.9"
http = "0.2.3"
tokio = { version = "1.23.1", features = ["macros", "rt", "rt-multi-thread", "test-util", "time"] }
tracing-subscriber = { version = "0.3.15", features = ["env-filter"] }
proptest = "1"
Expand Down
3 changes: 2 additions & 1 deletion aws/rust-runtime/aws-runtime/external-types.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
allowed_external_types = [
"aws_credential_types::*",
"aws_sigv4::*",
"aws_smithy_http::body::SdkBody",
"aws_smithy_http::*",
"aws_smithy_types::*",
"aws_smithy_runtime_api::*",
"aws_types::*",
"http::request::Request",
Expand Down
3 changes: 3 additions & 0 deletions aws/rust-runtime/aws-runtime/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,5 +25,8 @@ pub mod recursion_detection;
/// Supporting code for user agent headers in the AWS SDK.
pub mod user_agent;

/// Supporting code for retry behavior specific to the AWS SDK.
pub mod retries;

/// Supporting code for invocation ID headers in the AWS SDK.
pub mod invocation_id;
7 changes: 7 additions & 0 deletions aws/rust-runtime/aws-runtime/src/retries.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

/// Classifiers that can inspect a response and determine if it should be retried.
pub mod classifier;
174 changes: 174 additions & 0 deletions aws/rust-runtime/aws-runtime/src/retries/classifier.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,174 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

use aws_smithy_http::http::HttpHeaders;
use aws_smithy_http::result::SdkError;
use aws_smithy_runtime_api::client::retries::RetryReason;
use aws_smithy_types::error::metadata::ProvideErrorMetadata;
use aws_smithy_types::retry::ErrorKind;

/// AWS error codes that represent throttling errors.
pub const THROTTLING_ERRORS: &[&str] = &[
"Throttling",
"ThrottlingException",
"ThrottledException",
"RequestThrottledException",
"TooManyRequestsException",
"ProvisionedThroughputExceededException",
"TransactionInProgressException",
"RequestLimitExceeded",
"BandwidthLimitExceeded",
"LimitExceededException",
"RequestThrottled",
"SlowDown",
"PriorRequestNotComplete",
"EC2ThrottledException",
];

/// AWS error codes that represent transient errors.
pub const TRANSIENT_ERRORS: &[&str] = &["RequestTimeout", "RequestTimeoutException"];

/// A retry classifier for determining if the response sent by an AWS service requires a retry.
#[derive(Debug)]
pub struct AwsErrorCodeClassifier;

impl AwsErrorCodeClassifier {
/// Classify an error code to check if represents a retryable error. The codes of retryable
/// errors are defined [here](THROTTLING_ERRORS) and [here](TRANSIENT_ERRORS).
pub fn classify_error<E: ProvideErrorMetadata, R>(
&self,
error: &SdkError<E, R>,
) -> Option<RetryReason> {
if let Some(error_code) = error.code() {
if THROTTLING_ERRORS.contains(&error_code) {
return Some(RetryReason::Error(ErrorKind::ThrottlingError));
} else if TRANSIENT_ERRORS.contains(&error_code) {
return Some(RetryReason::Error(ErrorKind::TransientError));
}
};

None
}
}

/// A retry classifier that checks for `x-amz-retry-after` headers. If one is found, a
/// [`RetryReason::Explicit`] is returned containing the duration to wait before retrying.
#[derive(Debug)]
pub struct AmzRetryAfterHeaderClassifier;

impl AmzRetryAfterHeaderClassifier {
/// Classify an AWS responses error code to determine how (and if) it should be retried.
pub fn classify_error<E>(&self, error: &SdkError<E>) -> Option<RetryReason> {
error
.raw_response()
.and_then(|res| res.http_headers().get("x-amz-retry-after"))
.and_then(|header| header.to_str().ok())
.and_then(|header| header.parse::<u64>().ok())
.map(|retry_after_delay| {
RetryReason::Explicit(std::time::Duration::from_millis(retry_after_delay))
})
}
}

#[cfg(test)]
mod test {
use super::{AmzRetryAfterHeaderClassifier, AwsErrorCodeClassifier};
use aws_smithy_http::body::SdkBody;
use aws_smithy_http::operation;
use aws_smithy_http::result::SdkError;
use aws_smithy_runtime_api::client::retries::RetryReason;
use aws_smithy_types::error::metadata::ProvideErrorMetadata;
use aws_smithy_types::error::ErrorMetadata;
use aws_smithy_types::retry::{ErrorKind, ProvideErrorKind};
use std::fmt;
use std::time::Duration;

#[derive(Debug)]
struct UnmodeledError;

impl fmt::Display for UnmodeledError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "UnmodeledError")
}
}

impl std::error::Error for UnmodeledError {}

struct CodedError {
metadata: ErrorMetadata,
}

impl CodedError {
fn new(code: &'static str) -> Self {
Self {
metadata: ErrorMetadata::builder().code(code).build(),
}
}
}

impl ProvideErrorKind for UnmodeledError {
fn retryable_error_kind(&self) -> Option<ErrorKind> {
None
}

fn code(&self) -> Option<&str> {
None
}
}

impl ProvideErrorMetadata for CodedError {
fn meta(&self) -> &ErrorMetadata {
&self.metadata
}
}

#[test]
fn classify_by_error_code() {
let policy = AwsErrorCodeClassifier;
let res = http::Response::new("OK");
let err = SdkError::service_error(CodedError::new("Throttling"), res);

assert_eq!(
policy.classify_error(&err),
Some(RetryReason::Error(ErrorKind::ThrottlingError))
);

let res = http::Response::new("OK");
let err = SdkError::service_error(CodedError::new("RequestTimeout"), res);
assert_eq!(
policy.classify_error(&err),
Some(RetryReason::Error(ErrorKind::TransientError))
)
}

#[test]
fn classify_generic() {
let policy = AwsErrorCodeClassifier;
let res = http::Response::new("OK");
let err = aws_smithy_types::Error::builder().code("SlowDown").build();
let err = SdkError::service_error(err, res);
assert_eq!(
policy.classify_error(&err),
Some(RetryReason::Error(ErrorKind::ThrottlingError))
);
}

#[test]
fn test_retry_after_header() {
let policy = AmzRetryAfterHeaderClassifier;
let res = http::Response::builder()
.header("x-amz-retry-after", "5000")
.body("retry later")
.unwrap()
.map(SdkBody::from);
let res = operation::Response::new(res);
let err = SdkError::service_error(UnmodeledError, res);

assert_eq!(
policy.classify_error(&err),
Some(RetryReason::Explicit(Duration::from_millis(5000))),
);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,12 @@ package software.amazon.smithy.rustsdk
import software.amazon.smithy.model.shapes.OperationShape
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationRuntimePluginCustomization
import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationRuntimePluginSection
import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
import software.amazon.smithy.rust.codegen.core.rustlang.Attribute.Companion.derive
import software.amazon.smithy.rust.codegen.core.rustlang.rust
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
import software.amazon.smithy.rust.codegen.core.rustlang.writable
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
Expand All @@ -23,13 +28,21 @@ class RetryClassifierDecorator : ClientCodegenDecorator {
codegenContext: ClientCodegenContext,
operation: OperationShape,
baseCustomizations: List<OperationCustomization>,
): List<OperationCustomization> {
return baseCustomizations + RetryClassifierFeature(codegenContext.runtimeConfig)
}
): List<OperationCustomization> =
baseCustomizations + RetryClassifierFeature(codegenContext.runtimeConfig)

override fun operationRuntimePluginCustomizations(
codegenContext: ClientCodegenContext,
operation: OperationShape,
baseCustomizations: List<OperationRuntimePluginCustomization>,
): List<OperationRuntimePluginCustomization> =
baseCustomizations + OperationRetryClassifiersFeature(codegenContext, operation)
}

class RetryClassifierFeature(private val runtimeConfig: RuntimeConfig) : OperationCustomization() {
override fun retryType(): RuntimeType = AwsRuntimeType.awsHttp(runtimeConfig).resolve("retry::AwsResponseRetryClassifier")
override fun retryType(): RuntimeType =
AwsRuntimeType.awsHttp(runtimeConfig).resolve("retry::AwsResponseRetryClassifier")

override fun section(section: OperationSection) = when (section) {
is OperationSection.FinalizeOperation -> writable {
rust(
Expand All @@ -41,3 +54,140 @@ class RetryClassifierFeature(private val runtimeConfig: RuntimeConfig) : Operati
else -> emptySection
}
}

class OperationRetryClassifiersFeature(
codegenContext: ClientCodegenContext,
operation: OperationShape,
) : OperationRuntimePluginCustomization() {
private val runtimeConfig = codegenContext.runtimeConfig
private val awsRuntime = AwsRuntimeType.awsRuntime(runtimeConfig)
private val smithyRuntime = RuntimeType.smithyRuntime(runtimeConfig)
private val smithyRuntimeApi = RuntimeType.smithyRuntimeApi(runtimeConfig)
private val codegenScope = arrayOf(
"HttpStatusCodeClassifier" to smithyRuntime.resolve("client::retries::classifier::HttpStatusCodeClassifier"),
"AwsErrorCodeClassifier" to awsRuntime.resolve("retries::classifier::AwsErrorCodeClassifier"),
"ModeledAsRetryableClassifier" to smithyRuntime.resolve("client::retries::classifier::ModeledAsRetryableClassifier"),
"AmzRetryAfterHeaderClassifier" to awsRuntime.resolve("retries::classifier::AmzRetryAfterHeaderClassifier"),
"SmithyErrorClassifier" to smithyRuntime.resolve("client::retries::classifier::SmithyErrorClassifier"),
"RetryReason" to smithyRuntimeApi.resolve("client::retries::RetryReason"),
"ClassifyRetry" to smithyRuntimeApi.resolve("client::retries::ClassifyRetry"),
"RetryClassifiers" to smithyRuntimeApi.resolve("client::retries::RetryClassifiers"),
"OperationError" to codegenContext.symbolProvider.symbolForOperationError(operation),
"SdkError" to RuntimeType.smithyHttp(runtimeConfig).resolve("result::SdkError"),
"ErasedError" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("type_erasure::TypeErasedBox"),
)

override fun section(section: OperationRuntimePluginSection) = when (section) {
is OperationRuntimePluginSection.RuntimePluginSupportingTypes -> writable {
Attribute(derive(RuntimeType.Debug)).render(this)
rustTemplate(
"""
struct HttpStatusCodeClassifier(#{HttpStatusCodeClassifier});
impl HttpStatusCodeClassifier {
fn new() -> Self {
Self(#{HttpStatusCodeClassifier}::default())
}
}
impl #{ClassifyRetry} for HttpStatusCodeClassifier {
fn classify_retry(&self, error: &#{ErasedError}) -> Option<#{RetryReason}> {
let error = error.downcast_ref::<#{SdkError}<#{OperationError}>>().expect("The error type is always known");
self.0.classify_error(error)
}
}
""",
*codegenScope,
)

Attribute(derive(RuntimeType.Debug)).render(this)
rustTemplate(
"""
struct AwsErrorCodeClassifier(#{AwsErrorCodeClassifier});
impl AwsErrorCodeClassifier {
fn new() -> Self {
Self(#{AwsErrorCodeClassifier})
}
}
impl #{ClassifyRetry} for AwsErrorCodeClassifier {
fn classify_retry(&self, error: &#{ErasedError}) -> Option<#{RetryReason}> {
let error = error.downcast_ref::<#{SdkError}<#{OperationError}>>().expect("The error type is always known");
self.0.classify_error(error)
}
}
""",
*codegenScope,
)

Attribute(derive(RuntimeType.Debug)).render(this)
rustTemplate(
"""
struct ModeledAsRetryableClassifier(#{ModeledAsRetryableClassifier});
impl ModeledAsRetryableClassifier {
fn new() -> Self {
Self(#{ModeledAsRetryableClassifier})
}
}
impl #{ClassifyRetry} for ModeledAsRetryableClassifier {
fn classify_retry(&self, error: &#{ErasedError}) -> Option<#{RetryReason}> {
let error = error.downcast_ref::<#{SdkError}<#{OperationError}>>().expect("The error type is always known");
self.0.classify_error(error)
}
}
""",
*codegenScope,
)

Attribute(derive(RuntimeType.Debug)).render(this)
rustTemplate(
"""
struct AmzRetryAfterHeaderClassifier(#{AmzRetryAfterHeaderClassifier});
impl AmzRetryAfterHeaderClassifier {
fn new() -> Self {
Self(#{AmzRetryAfterHeaderClassifier})
}
}
impl #{ClassifyRetry} for AmzRetryAfterHeaderClassifier {
fn classify_retry(&self, error: &#{ErasedError}) -> Option<#{RetryReason}> {
let error = error.downcast_ref::<#{SdkError}<#{OperationError}>>().expect("The error type is always known");
self.0.classify_error(error)
}
}
""",
*codegenScope,
)

Attribute(derive(RuntimeType.Debug)).render(this)
rustTemplate(
"""
struct SmithyErrorClassifier(#{SmithyErrorClassifier});
impl SmithyErrorClassifier {
fn new() -> Self {
Self(#{SmithyErrorClassifier})
}
}
impl #{ClassifyRetry} for SmithyErrorClassifier {
fn classify_retry(&self, error: &#{ErasedError}) -> Option<#{RetryReason}> {
let error = error.downcast_ref::<#{SdkError}<#{OperationError}>>().expect("The error type is always known");
self.0.classify_error(error)
}
}
""",
*codegenScope,
)
}

is OperationRuntimePluginSection.RetryClassifier -> writable {
rustTemplate(
"""
.with_classifier(SmithyErrorClassifier::new())
.with_classifier(AmzRetryAfterHeaderClassifier::new())
.with_classifier(ModeledAsRetryableClassifier::new())
.with_classifier(AwsErrorCodeClassifier::new())
.with_classifier(HttpStatusCodeClassifier::new())
""",
*codegenScope,
)
}

else -> emptySection
}
}
Loading

0 comments on commit ae995fb

Please sign in to comment.