Update retry classifiers to act on InterceptorContext instead of OrchestratorError (#2737)

## Description
<!--- Describe your changes in detail -->
- update retry classifiers to act directly on the `InterceptorContext`
- add codegen method for inserting retry classifiers into the config bag
- update `InterceptorContext` to return options instead of panicking

## 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 were updated as needed.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

Author: Zelda Hessler

aws/rust-runtime/aws-runtime/src/invocation_id.rs

@@ -80,7 +80,12 @@ mod tests {
     use http::HeaderValue;
 
     fn expect_header<'a>(context: &'a InterceptorContext, header_name: &str) -> &'a HeaderValue {
-        context.request().headers().get(header_name).unwrap()
+        context
+            .request()
+            .expect("request is set")
+            .headers()
+            .get(header_name)
+            .unwrap()
     }
 
     #[test]

aws/rust-runtime/aws-runtime/src/recursion_detection.rs

@@ -159,7 +159,7 @@ mod tests {
         RecursionDetectionInterceptor { env }
             .modify_before_signing(&mut ctx, &mut config)
             .expect("interceptor must succeed");
-        let mutated_request = context.request();
+        let mutated_request = context.request().expect("request is set");
         for name in mutated_request.headers().keys() {
             assert_eq!(
                 mutated_request.headers().get_all(name).iter().count(),

aws/rust-runtime/aws-runtime/src/request_info.rs

@@ -168,6 +168,7 @@ mod tests {
     fn expect_header<'a>(context: &'a InterceptorContext, header_name: &str) -> &'a str {
         context
             .request()
+            .expect("request is set")
             .headers()
             .get(header_name)
             .unwrap()

aws/rust-runtime/aws-runtime/src/retries/classifier.rs

@@ -4,10 +4,13 @@
  */
 
 use aws_smithy_http::http::HttpHeaders;
-use aws_smithy_http::result::SdkError;
-use aws_smithy_runtime_api::client::retries::RetryReason;
+use aws_smithy_runtime_api::client::interceptors::InterceptorContext;
+use aws_smithy_runtime_api::client::orchestrator::OrchestratorError;
+use aws_smithy_runtime_api::client::retries::{ClassifyRetry, RetryReason};
 use aws_smithy_types::error::metadata::ProvideErrorMetadata;
 use aws_smithy_types::retry::ErrorKind;
+use std::error::Error as StdError;
+use std::marker::PhantomData;
 
 /// AWS error codes that represent throttling errors.
 pub const THROTTLING_ERRORS: &[&str] = &[
@@ -31,16 +34,31 @@ pub const THROTTLING_ERRORS: &[&str] = &[
 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> {
+#[derive(Debug, Default)]
+pub struct AwsErrorCodeClassifier<E> {
+    _inner: PhantomData<E>,
+}
+
+impl<E> AwsErrorCodeClassifier<E> {
+    /// Create a new AwsErrorCodeClassifier
+    pub fn new() -> Self {
+        Self {
+            _inner: PhantomData,
+        }
+    }
+}
+
+impl<E> ClassifyRetry for AwsErrorCodeClassifier<E>
+where
+    E: StdError + ProvideErrorMetadata + Send + Sync + 'static,
+{
+    fn classify_retry(&self, ctx: &InterceptorContext) -> Option<RetryReason> {
+        let error = ctx
+            .output_or_error()?
+            .err()
+            .and_then(OrchestratorError::as_operation_error)?
+            .downcast_ref::<E>()?;
+
         if let Some(error_code) = error.code() {
             if THROTTLING_ERRORS.contains(&error_code) {
                 return Some(RetryReason::Error(ErrorKind::ThrottlingError));
@@ -51,37 +69,51 @@ impl AwsErrorCodeClassifier {
 
         None
     }
+
+    fn name(&self) -> &'static str {
+        "AWS Error Code"
+    }
 }
 
 /// 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)]
+#[derive(Debug, Default)]
 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()
+    /// Create a new `AmzRetryAfterHeaderClassifier`.
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+impl ClassifyRetry for AmzRetryAfterHeaderClassifier {
+    fn classify_retry(&self, ctx: &InterceptorContext) -> Option<RetryReason> {
+        ctx.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))
             })
     }
+
+    fn name(&self) -> &'static str {
+        "'Retry After' Header"
+    }
 }
 
 #[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_runtime_api::client::interceptors::InterceptorContext;
+    use aws_smithy_runtime_api::client::orchestrator::OrchestratorError;
+    use aws_smithy_runtime_api::client::retries::{ClassifyRetry, RetryReason};
     use aws_smithy_types::error::metadata::ProvideErrorMetadata;
     use aws_smithy_types::error::ErrorMetadata;
     use aws_smithy_types::retry::{ErrorKind, ProvideErrorKind};
+    use aws_smithy_types::type_erasure::{TypeErasedBox, TypeErasedError};
     use std::fmt;
     use std::time::Duration;
 
@@ -96,6 +128,17 @@ mod test {
 
     impl std::error::Error for UnmodeledError {}
 
+    impl ProvideErrorKind for UnmodeledError {
+        fn retryable_error_kind(&self) -> Option<ErrorKind> {
+            None
+        }
+
+        fn code(&self) -> Option<&str> {
+            None
+        }
+    }
+
+    #[derive(Debug)]
     struct CodedError {
         metadata: ErrorMetadata,
     }
@@ -108,16 +151,14 @@ mod test {
         }
     }
 
-    impl ProvideErrorKind for UnmodeledError {
-        fn retryable_error_kind(&self) -> Option<ErrorKind> {
-            None
-        }
-
-        fn code(&self) -> Option<&str> {
-            None
+    impl fmt::Display for CodedError {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            write!(f, "Coded Error")
         }
     }
 
+    impl std::error::Error for CodedError {}
+
     impl ProvideErrorMetadata for CodedError {
         fn meta(&self) -> &ErrorMetadata {
             &self.metadata
@@ -126,31 +167,39 @@ mod test {
 
     #[test]
     fn classify_by_error_code() {
-        let policy = AwsErrorCodeClassifier;
-        let res = http::Response::new("OK");
-        let err = SdkError::service_error(CodedError::new("Throttling"), res);
+        let policy = AwsErrorCodeClassifier::<CodedError>::new();
+        let mut ctx = InterceptorContext::new(TypeErasedBox::new("doesntmatter"));
+        ctx.set_output_or_error(Err(OrchestratorError::operation(TypeErasedError::new(
+            CodedError::new("Throttling"),
+        ))));
 
         assert_eq!(
-            policy.classify_error(&err),
+            policy.classify_retry(&ctx),
             Some(RetryReason::Error(ErrorKind::ThrottlingError))
         );
 
-        let res = http::Response::new("OK");
-        let err = SdkError::service_error(CodedError::new("RequestTimeout"), res);
+        let mut ctx = InterceptorContext::new(TypeErasedBox::new("doesntmatter"));
+        ctx.set_output_or_error(Err(OrchestratorError::operation(TypeErasedError::new(
+            CodedError::new("RequestTimeout"),
+        ))));
         assert_eq!(
-            policy.classify_error(&err),
+            policy.classify_retry(&ctx),
             Some(RetryReason::Error(ErrorKind::TransientError))
         )
     }
 
     #[test]
     fn classify_generic() {
-        let policy = AwsErrorCodeClassifier;
-        let res = http::Response::new("OK");
+        let policy = AwsErrorCodeClassifier::<ErrorMetadata>::new();
         let err = aws_smithy_types::Error::builder().code("SlowDown").build();
-        let err = SdkError::service_error(err, res);
+        let test_response = http::Response::new("OK").map(SdkBody::from);
+
+        let mut ctx = InterceptorContext::new(TypeErasedBox::new("doesntmatter"));
+        ctx.set_response(test_response);
+        ctx.set_output_or_error(Err(OrchestratorError::operation(TypeErasedError::new(err))));
+
         assert_eq!(
-            policy.classify_error(&err),
+            policy.classify_retry(&ctx),
             Some(RetryReason::Error(ErrorKind::ThrottlingError))
         );
     }
@@ -163,11 +212,14 @@ mod test {
             .body("retry later")
             .unwrap()
             .map(SdkBody::from);
-        let res = operation::Response::new(res);
-        let err = SdkError::service_error(UnmodeledError, res);
+        let mut ctx = InterceptorContext::new(TypeErasedBox::new("doesntmatter"));
+        ctx.set_response(res);
+        ctx.set_output_or_error(Err(OrchestratorError::operation(TypeErasedError::new(
+            UnmodeledError,
+        ))));
 
         assert_eq!(
-            policy.classify_error(&err),
+            policy.classify_retry(&ctx),
             Some(RetryReason::Explicit(Duration::from_millis(5000))),
         );
     }

aws/rust-runtime/aws-runtime/src/user_agent.rs

@@ -117,6 +117,7 @@ mod tests {
     fn expect_header<'a>(context: &'a InterceptorContext, header_name: &str) -> &'a str {
         context
             .request()
+            .expect("request is set")
             .headers()
             .get(header_name)
             .unwrap()