Add Connection Poisoning to aws-smithy-client

CHANGELOG.next.toml

@@ -263,3 +263,28 @@ message = "Add more client re-exports. Specifically, it re-exports `aws_smithy_h
 references = ["smithy-rs#2437", "aws-sdk-rust#600"]
 meta = { "breaking" = false, "tada" = false, "bug" = false, "target" = "client" }
 author = "ysaito1001"
+
+[[aws-sdk-rust]]
+message = """Reconnect on transient errors.
+
+If a transient error (timeout, 500, 503, 503) is encountered, the connection will be evicted from the pool and will not
+be reused. This is enabled by default for all AWS services. It can be disabled by setting `RetryConfig::with_reconnect_mode`
+
+Although there is no API breakage from this change, it alters the client behavior in a way that may cause breakage for customers.
+"""
+references = ["aws-sdk-rust#160", "smithy-rs#2445"]
+meta = { "breaking" = true, "tada" = false, "bug" = false }
+author = "rcoh"
+
+[[smithy-rs]]
+message = """Reconnect on transient errors.
+
+Note: **this behavior is disabled by default for generic clients**. It can be enabled with
+`aws_smithy_client::Builder::reconnect_on_transient_errors`
+
+If a transient error (timeout, 500, 503, 503) is encountered, the connection will be evicted from the pool and will not
+be reused.
+"""
+references = ["aws-sdk-rust#160", "smithy-rs#2445"]
+meta = { "breaking" = false, "tada" = false, "bug" = false }
+author = "rcoh"

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsFluentClientDecorator.kt

@@ -210,6 +210,7 @@ private class AwsFluentClientExtensions(types: Types) {
                     };
                     let mut builder = builder
                         .middleware(#{DynMiddleware}::new(#{Middleware}::new()))
+                        .reconnect_mode(retry_config.reconnect_mode())
                         .retry_config(retry_config.into())
                         .operation_timeout_config(timeout_config.into());
                     builder.set_sleep_impl(sleep_impl);

aws/sdk/build.gradle.kts

@@ -302,6 +302,9 @@ fun generateCargoWorkspace(services: AwsServices): String {
     |]
     |members = [${"\n"}${services.allModules.joinToString(",\n") { "|    \"$it\"" }}
     |]
+    |
+    |[patch.crates-io]
+    |hyper = { git = 'https://github.com/hyperium/hyper', branch = "0.14.x" }
     """.trimMargin()
 }
 

aws/sdk/integration-tests/Cargo.toml

@@ -17,3 +17,6 @@ members = [
     "transcribestreaming",
     "using-native-tls-instead-of-rustls",
 ]
+
+[patch.crates-io]
+hyper = { git = 'https://github.com/hyperium/hyper', branch = "0.14.x" }

aws/sdk/integration-tests/s3/tests/reconnects.rs

@@ -0,0 +1,62 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_credential_types::provider::SharedCredentialsProvider;
+use aws_credential_types::Credentials;
+use aws_smithy_async::rt::sleep::TokioSleep;
+use aws_smithy_client::test_connection::wire_mock::{
+    check_matches, ReplayedEvent, WireLevelTestConnection,
+};
+use aws_smithy_client::{ev, match_events};
+use aws_smithy_types::retry::RetryConfig;
+use aws_types::region::Region;
+use aws_types::SdkConfig;
+use std::sync::Arc;
+use tracing_subscriber::EnvFilter;
+
+#[tokio::test]
+async fn reconnect_on_503() {
+    tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::new("trace"))
+        .init();
+    let mock = WireLevelTestConnection::spinup(vec![
+        ReplayedEvent::status(503),
+        ReplayedEvent::status(503),
+        ReplayedEvent::with_body("here-is-your-object"),
+    ])
+    .await;
+
+    let sdk_config = SdkConfig::builder()
+        .region(Region::from_static("us-east-2"))
+        .credentials_provider(SharedCredentialsProvider::new(Credentials::for_tests()))
+        .sleep_impl(Arc::new(TokioSleep::new()))
+        .endpoint_url(mock.endpoint_url())
+        .http_connector(mock.http_connector())
+        .retry_config(RetryConfig::standard())
+        .build();
+    let client = aws_sdk_s3::Client::new(&sdk_config);
+    let resp = client
+        .get_object()
+        .bucket("bucket")
+        .key("key")
+        .send()
+        .await
+        .expect("succeeds after retries");
+    assert_eq!(
+        resp.body.collect().await.unwrap().to_vec(),
+        b"here-is-your-object"
+    );
+    match_events!(
+        ev!(dns),
+        ev!(connect),
+        ev!(http(503)),
+        ev!(dns),
+        ev!(connect),
+        ev!(http(503)),
+        ev!(dns),
+        ev!(connect),
+        ev!(http(200))
+    )(&mock.events());
+}

rust-runtime/Cargo.toml

@@ -1,5 +1,6 @@
 [workspace]
 
+
 members = [
     "inlineable",
     "aws-smithy-async",
@@ -18,3 +19,6 @@ members = [
     "aws-smithy-http-server",
     "aws-smithy-http-server-python",
 ]
+
+[patch.crates-io]
+hyper = { git = 'https://github.com/hyperium/hyper', branch = "0.14.x" }

rust-runtime/aws-smithy-client/Cargo.toml

@@ -9,12 +9,13 @@ repository = "https://github.com/awslabs/smithy-rs"
 
 [features]
 rt-tokio = ["aws-smithy-async/rt-tokio"]
-test-util = ["aws-smithy-protocol-test", "serde/derive", "rustls"]
+test-util = ["aws-smithy-protocol-test", "serde/derive", "rustls", "hyper/server", "hyper/h2"]
 native-tls = ["client-hyper", "hyper-tls", "rt-tokio"]
 rustls = ["client-hyper", "hyper-rustls", "rt-tokio", "lazy_static"]
 client-hyper = ["hyper"]
 hyper-webpki-doctest-only = ["hyper-rustls/webpki-roots"]
 
+
 [dependencies]
 aws-smithy-async = { path = "../aws-smithy-async" }
 aws-smithy-http = { path = "../aws-smithy-http" }
@@ -25,7 +26,7 @@ bytes = "1"
 fastrand = "1.4.0"
 http = "0.2.3"
 http-body = "0.4.4"
-hyper = { version = "0.14.12", features = ["client", "http2", "http1", "tcp"], optional = true }
+hyper = { version = "0.14.24", features = ["client", "http2", "http1", "tcp"], optional = true }
 # cargo does not support optional test dependencies, so to completely disable rustls when
 # the native-tls feature is enabled, we need to add the webpki-roots feature here.
 # https://github.com/rust-lang/cargo/issues/1596
@@ -44,6 +45,9 @@ serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 tokio = { version = "1.8.4", features = ["full", "test-util"] }
 tower-test = "0.4.0"
+tracing-subscriber = "0.3.16"
+tracing-test = "0.2.4"
+
 
 [package.metadata.docs.rs]
 all-features = true

rust-runtime/aws-smithy-client/src/builder.rs

@@ -7,6 +7,7 @@ use crate::{bounds, erase, retry, Client};
 use aws_smithy_async::rt::sleep::{default_async_sleep, AsyncSleep};
 use aws_smithy_http::body::SdkBody;
 use aws_smithy_http::result::ConnectorError;
+use aws_smithy_types::retry::ReconnectMode;
 use aws_smithy_types::timeout::{OperationTimeoutConfig, TimeoutConfig};
 use std::sync::Arc;
 
@@ -37,6 +38,12 @@ pub struct Builder<C = (), M = (), R = retry::Standard> {
     retry_policy: MaybeRequiresSleep<R>,
     operation_timeout_config: Option<OperationTimeoutConfig>,
     sleep_impl: Option<Arc<dyn AsyncSleep>>,
+    reconnect_mode: Option<ReconnectMode>,
+}
+
+/// transitional default: disable this behavior by default
+fn default_reconnect_mode() -> ReconnectMode {
+    ReconnectMode::NoReconnect
 }
 
 impl<C, M> Default for Builder<C, M>
@@ -55,6 +62,7 @@ where
             ),
             operation_timeout_config: None,
             sleep_impl: default_async_sleep(),
+            reconnect_mode: Some(default_reconnect_mode()),
         }
     }
 }
@@ -173,6 +181,7 @@ impl<M, R> Builder<(), M, R> {
             retry_policy: self.retry_policy,
             operation_timeout_config: self.operation_timeout_config,
             sleep_impl: self.sleep_impl,
+            reconnect_mode: self.reconnect_mode,
         }
     }
 
@@ -229,6 +238,7 @@ impl<C, R> Builder<C, (), R> {
             operation_timeout_config: self.operation_timeout_config,
             middleware,
             sleep_impl: self.sleep_impl,
+            reconnect_mode: self.reconnect_mode,
         }
     }
 
@@ -280,6 +290,7 @@ impl<C, M> Builder<C, M, retry::Standard> {
             operation_timeout_config: self.operation_timeout_config,
             middleware: self.middleware,
             sleep_impl: self.sleep_impl,
+            reconnect_mode: self.reconnect_mode,
         }
     }
 }
@@ -347,6 +358,7 @@ impl<C, M, R> Builder<C, M, R> {
             retry_policy: self.retry_policy,
             operation_timeout_config: self.operation_timeout_config,
             sleep_impl: self.sleep_impl,
+            reconnect_mode: self.reconnect_mode,
         }
     }
 
@@ -361,9 +373,41 @@ impl<C, M, R> Builder<C, M, R> {
             retry_policy: self.retry_policy,
             operation_timeout_config: self.operation_timeout_config,
             sleep_impl: self.sleep_impl,
+            reconnect_mode: self.reconnect_mode,
         }
     }
 
+    /// Set the [`ReconnectMode`] for the retry strategy
+    ///
+    /// By default, no reconnection occurs.
+    ///
+    /// When enabled and a transient error is encountered, the connection in use will be poisoned.
+    /// This prevents reusing a connection to a potentially bad host.
+    pub fn reconnect_mode(mut self, reconnect_mode: ReconnectMode) -> Self {
+        self.set_reconnect_mode(Some(reconnect_mode));
+        self
+    }
+
+    /// Set the [`ReconnectMode`] for the retry strategy
+    ///
+    /// By default, no reconnection occurs.
+    ///
+    /// When enabled and a transient error is encountered, the connection in use will be poisoned.
+    /// This prevents reusing a connection to a potentially bad host.
+    pub fn set_reconnect_mode(&mut self, reconnect_mode: Option<ReconnectMode>) -> &mut Self {
+        self.reconnect_mode = reconnect_mode;
+        self
+    }
+
+    /// Enable reconnection on transient errors
+    ///
+    /// By default, when a transient error is encountered, the connection in use will be poisoned.
+    /// This prevents reusing a connection to a potentially bad host but may increase the load on
+    /// the server.
+    pub fn reconnect_on_transient_errors(self) -> Self {
+        self.reconnect_mode(ReconnectMode::ReconnectOnTransientError)
+    }
+
     /// Build a Smithy service [`Client`].
     pub fn build(self) -> Client<C, M, R> {
         let operation_timeout_config = self
@@ -392,6 +436,7 @@ impl<C, M, R> Builder<C, M, R> {
             middleware: self.middleware,
             operation_timeout_config,
             sleep_impl: self.sleep_impl,
+            reconnect_mode: self.reconnect_mode.unwrap_or(ReconnectMode::NoReconnect),
         }
     }
 }

rust-runtime/aws-smithy-client/src/erase.rs

@@ -61,6 +61,7 @@ where
             retry_policy: self.retry_policy,
             operation_timeout_config: self.operation_timeout_config,
             sleep_impl: self.sleep_impl,
+            reconnect_mode: self.reconnect_mode,
         }
     }
 }
@@ -101,6 +102,7 @@ where
             retry_policy: self.retry_policy,
             operation_timeout_config: self.operation_timeout_config,
             sleep_impl: self.sleep_impl,
+            reconnect_mode: self.reconnect_mode,
         }
     }