add connection poisoning support to the orchestrator (#2824)

This PR updates the orchestrator to support connection poisoning for the
`hyper` connector. It also renames many usages of "connection" to
"connector" in order to make talking about these concepts less
confusing. In short:

```
/// A "connector" manages one or more "connections", handles connection timeouts, re-establishes
/// connections, etc.
///
/// "Connections" refers to the actual transport layer implementation of the connector.
/// By default, the orchestrator uses a connector provided by `hyper`.
```

One possibly controversial decision I made is that `reconnect_mode` is
now a standalone configurable field, rather that a sub-field of
`RetryConfig`. I think we should get together as a team and discuss what
kinds of things we want to allow users to do when configuring
connections.

My thoughts on this are that we should either:

- **A.** Make connection-related settings their own types instead of
including them within other config types
- **B.** Make a single config struct for all connector-related stuff/use
the preëxisting `ConnectorSettings` struct.

Personally, I'm partial to A.

----

_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>

Author: Zelda Hessler

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

@@ -15,6 +15,7 @@ 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
+import software.amazon.smithy.rust.codegen.core.util.letIf
 
 class RetryClassifierDecorator : ClientCodegenDecorator {
     override val name: String = "RetryPolicy"
@@ -25,10 +26,12 @@ class RetryClassifierDecorator : ClientCodegenDecorator {
         operation: OperationShape,
         baseCustomizations: List<OperationCustomization>,
     ): List<OperationCustomization> =
-        baseCustomizations + RetryClassifierFeature(codegenContext.runtimeConfig) + OperationRetryClassifiersFeature(
-            codegenContext,
-            operation,
-        )
+        (baseCustomizations + RetryClassifierFeature(codegenContext.runtimeConfig)).letIf(codegenContext.smithyRuntimeMode.generateOrchestrator) {
+            it + OperationRetryClassifiersFeature(
+                codegenContext,
+                operation,
+            )
+        }
 }
 
 class RetryClassifierFeature(private val runtimeConfig: RuntimeConfig) : OperationCustomization() {

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/timestream/TimestreamDecorator.kt

@@ -52,7 +52,6 @@ class TimestreamDecorator : ClientCodegenDecorator {
             Visibility.PUBLIC,
             CargoDependency.Tokio.copy(scope = DependencyScope.Compile, features = setOf("sync")),
         )
-        val runtimeMode = codegenContext.smithyRuntimeMode
         rustCrate.lib {
             // helper function to resolve an endpoint given a base client
             rustTemplate(

aws/sdk/integration-tests/s3/Cargo.toml

@@ -12,8 +12,8 @@ publish = false
 
 [dev-dependencies]
 async-std = "1.12.0"
-aws-credential-types = { path = "../../build/aws-sdk/sdk/aws-credential-types", features = ["test-util"] }
 aws-config = { path = "../../build/aws-sdk/sdk/aws-config" }
+aws-credential-types = { path = "../../build/aws-sdk/sdk/aws-credential-types", features = ["test-util"] }
 aws-http = { path = "../../build/aws-sdk/sdk/aws-http" }
 aws-sdk-s3 = { path = "../../build/aws-sdk/sdk/s3" }
 aws-sdk-sts = { path = "../../build/aws-sdk/sdk/sts" }
@@ -36,9 +36,9 @@ serde_json = "1"
 smol = "1.2"
 tempfile = "3"
 tokio = { version = "1.23.1", features = ["macros", "test-util", "rt-multi-thread"] }
-# If you're writing a test with this, take heed! `no-env-filter` means you'll be capturing
-# logs from everything that speaks, so be specific with your asserts.
-tracing-test = { version = "0.2.4", features = ["no-env-filter"] }
 tracing = "0.1.37"
 tracing-appender = "0.2.2"
 tracing-subscriber = { version = "0.3.15", features = ["env-filter", "json"] }
+# If you're writing a test with this, take heed! `no-env-filter` means you'll be capturing
+# logs from everything that speaks, so be specific with your asserts.
+tracing-test = { version = "0.2.4", features = ["no-env-filter"] }

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

@@ -3,38 +3,34 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_credential_types::Credentials;
-use aws_smithy_async::rt::sleep::{SharedAsyncSleep, TokioSleep};
+use aws_sdk_s3::config::retry::{ReconnectMode, RetryConfig};
+use aws_sdk_s3::config::{Credentials, Region, SharedAsyncSleep};
+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::{ReconnectMode, RetryConfig};
-use aws_types::region::Region;
-use aws_types::SdkConfig;
 
 #[tokio::test]
-/// test that disabling reconnects on retry config disables them for the client
-async fn disable_reconnects() {
+async fn test_disable_reconnect_on_503() {
     let mock = WireLevelTestConnection::spinup(vec![
         ReplayedEvent::status(503),
         ReplayedEvent::status(503),
         ReplayedEvent::with_body("here-is-your-object"),
     ])
     .await;
 
-    let sdk_config = SdkConfig::builder()
+    let config = aws_sdk_s3::Config::builder()
         .region(Region::from_static("us-east-2"))
-        .credentials_provider(SharedCredentialsProvider::new(Credentials::for_tests()))
+        .credentials_provider(Credentials::for_tests())
         .sleep_impl(SharedAsyncSleep::new(TokioSleep::new()))
         .endpoint_url(mock.endpoint_url())
         .http_connector(mock.http_connector())
         .retry_config(
             RetryConfig::standard().with_reconnect_mode(ReconnectMode::ReuseAllConnections),
         )
         .build();
-    let client = aws_sdk_s3::Client::new(&sdk_config);
+    let client = aws_sdk_s3::Client::from_conf(config);
     let resp = client
         .get_object()
         .bucket("bucket")
@@ -56,23 +52,25 @@ async fn disable_reconnects() {
 }
 
 #[tokio::test]
-async fn reconnect_on_503() {
+async fn test_enabling_reconnect_on_503() {
     let mock = WireLevelTestConnection::spinup(vec![
         ReplayedEvent::status(503),
         ReplayedEvent::status(503),
         ReplayedEvent::with_body("here-is-your-object"),
     ])
     .await;
 
-    let sdk_config = SdkConfig::builder()
+    let config = aws_sdk_s3::Config::builder()
         .region(Region::from_static("us-east-2"))
-        .credentials_provider(SharedCredentialsProvider::new(Credentials::for_tests()))
+        .credentials_provider(Credentials::for_tests())
         .sleep_impl(SharedAsyncSleep::new(TokioSleep::new()))
         .endpoint_url(mock.endpoint_url())
         .http_connector(mock.http_connector())
-        .retry_config(RetryConfig::standard())
+        .retry_config(
+            RetryConfig::standard().with_reconnect_mode(ReconnectMode::ReconnectOnTransientError),
+        )
         .build();
-    let client = aws_sdk_s3::Client::new(&sdk_config);
+    let client = aws_sdk_s3::Client::from_conf(config);
     let resp = client
         .get_object()
         .bucket("bucket")

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ConnectionPoisoningConfigCustomization.kt

@@ -0,0 +1,37 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.client.smithy.customizations
+
+import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginCustomization
+import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginSection
+import software.amazon.smithy.rust.codegen.core.rustlang.Writable
+import software.amazon.smithy.rust.codegen.core.rustlang.rust
+import software.amazon.smithy.rust.codegen.core.rustlang.writable
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.smithyRuntime
+
+class ConnectionPoisoningRuntimePluginCustomization(
+    codegenContext: ClientCodegenContext,
+) : ServiceRuntimePluginCustomization() {
+    private val runtimeConfig = codegenContext.runtimeConfig
+
+    override fun section(section: ServiceRuntimePluginSection): Writable = writable {
+        when (section) {
+            is ServiceRuntimePluginSection.RegisterInterceptor -> {
+                // This interceptor assumes that a compatible Connector is set. Otherwise, connection poisoning
+                // won't work and an error message will be logged.
+                section.registerInterceptor(runtimeConfig, this) {
+                    rust(
+                        "#T::new()",
+                        smithyRuntime(runtimeConfig).resolve("client::connectors::connection_poisoning::ConnectionPoisoningInterceptor"),
+                    )
+                }
+            }
+
+            else -> emptySection
+        }
+    }
+}

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/HttpConnectorConfigDecorator.kt

@@ -40,8 +40,8 @@ private class HttpConnectorConfigCustomization(
         *preludeScope,
         "Connection" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::orchestrator::Connection"),
         "ConnectorSettings" to RuntimeType.smithyClient(runtimeConfig).resolve("http_connector::ConnectorSettings"),
-        "DynConnection" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::orchestrator::DynConnection"),
-        "DynConnectorAdapter" to RuntimeType.smithyRuntime(runtimeConfig).resolve("client::connections::adapter::DynConnectorAdapter"),
+        "DynConnector" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::connectors::DynConnector"),
+        "DynConnectorAdapter" to RuntimeType.smithyRuntime(runtimeConfig).resolve("client::connectors::adapter::DynConnectorAdapter"),
         "HttpConnector" to RuntimeType.smithyClient(runtimeConfig).resolve("http_connector::HttpConnector"),
         "SharedAsyncSleep" to RuntimeType.smithyAsync(runtimeConfig).resolve("rt::sleep::SharedAsyncSleep"),
         "TimeoutConfig" to RuntimeType.smithyTypes(runtimeConfig).resolve("timeout::TimeoutConfig"),
@@ -197,14 +197,14 @@ private class HttpConnectorConfigCustomization(
 
                         let connector_settings = #{ConnectorSettings}::from_timeout_config(&timeout_config);
 
-                        if let Some(connection) = layer.load::<#{HttpConnector}>()
+                        if let Some(connector) = layer.load::<#{HttpConnector}>()
                                 .and_then(|c| c.connector(&connector_settings, sleep_impl.clone()))
                                 .or_else(|| #{default_connector}(&connector_settings, sleep_impl)) {
-                            let connection: #{DynConnection} = #{DynConnection}::new(#{DynConnectorAdapter}::new(
+                            let connector: #{DynConnector} = #{DynConnector}::new(#{DynConnectorAdapter}::new(
                                 // TODO(enableNewSmithyRuntimeCleanup): Replace the tower-based DynConnector and remove DynConnectorAdapter when deleting the middleware implementation
-                                connection
+                                connector
                             ));
-                            #{ConfigBagAccessors}::set_connection(&mut layer, connection);
+                            #{ConfigBagAccessors}::set_connector(&mut layer, connector);
                         }
 
                         """,

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ResiliencyConfigCustomization.kt

@@ -15,7 +15,6 @@ import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 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
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
@@ -26,8 +25,7 @@ class ResiliencyConfigCustomization(private val codegenContext: ClientCodegenCon
     private val retryConfig = RuntimeType.smithyTypes(runtimeConfig).resolve("retry")
     private val sleepModule = RuntimeType.smithyAsync(runtimeConfig).resolve("rt::sleep")
     private val timeoutModule = RuntimeType.smithyTypes(runtimeConfig).resolve("timeout")
-    private val smithyRuntimeCrate = RuntimeType.smithyRuntime(runtimeConfig)
-    private val retries = smithyRuntimeCrate.resolve("client::retries")
+    private val retries = RuntimeType.smithyRuntime(runtimeConfig).resolve("client::retries")
     private val moduleUseName = codegenContext.moduleUseName()
     private val codegenScope = arrayOf(
         *preludeScope,
@@ -367,7 +365,7 @@ class ResiliencyConfigCustomization(private val codegenContext: ClientCodegenCon
                     }
                 }
 
-                ServiceConfig.BuilderBuild -> {
+                is ServiceConfig.BuilderBuild -> {
                     if (runtimeMode.defaultToOrchestrator) {
                         rustTemplate(
                             """
@@ -420,7 +418,10 @@ class ResiliencyConfigCustomization(private val codegenContext: ClientCodegenCon
         }
 }
 
-class ResiliencyReExportCustomization(private val runtimeConfig: RuntimeConfig) {
+class ResiliencyReExportCustomization(codegenContext: ClientCodegenContext) {
+    private val runtimeConfig = codegenContext.runtimeConfig
+    private val runtimeMode = codegenContext.smithyRuntimeMode
+
     fun extras(rustCrate: RustCrate) {
         rustCrate.withModule(ClientRustModule.config) {
             rustTemplate(
@@ -430,13 +431,16 @@ class ResiliencyReExportCustomization(private val runtimeConfig: RuntimeConfig)
         }
         rustCrate.withModule(ClientRustModule.Config.retry) {
             rustTemplate(
-                "pub use #{types_retry}::{RetryConfig, RetryConfigBuilder, RetryMode};",
+                "pub use #{types_retry}::{RetryConfig, RetryConfigBuilder, RetryMode, ReconnectMode};",
                 "types_retry" to RuntimeType.smithyTypes(runtimeConfig).resolve("retry"),
             )
-            rustTemplate(
-                "pub use #{RetryPartition};",
-                "RetryPartition" to RuntimeType.smithyRuntime(runtimeConfig).resolve("client::retries::RetryPartition"),
-            )
+
+            if (runtimeMode.generateOrchestrator) {
+                rustTemplate(
+                    "pub use #{types_retry}::RetryPartition;",
+                    "types_retry" to RuntimeType.smithyRuntime(runtimeConfig).resolve("client::retries"),
+                )
+            }
         }
         rustCrate.withModule(ClientRustModule.Config.timeout) {
             rustTemplate(
@@ -449,30 +453,34 @@ class ResiliencyReExportCustomization(private val runtimeConfig: RuntimeConfig)
 
 class ResiliencyServiceRuntimePluginCustomization(codegenContext: ClientCodegenContext) : ServiceRuntimePluginCustomization() {
     private val runtimeConfig = codegenContext.runtimeConfig
-    private val smithyRuntimeCrate = RuntimeType.smithyRuntime(runtimeConfig)
-    private val retries = smithyRuntimeCrate.resolve("client::retries")
+    private val runtimeMode = codegenContext.smithyRuntimeMode
+    private val smithyRuntime = RuntimeType.smithyRuntime(runtimeConfig)
+    private val retries = smithyRuntime.resolve("client::retries")
     private val codegenScope = arrayOf(
         "TokenBucket" to retries.resolve("TokenBucket"),
         "TokenBucketPartition" to retries.resolve("TokenBucketPartition"),
         "ClientRateLimiter" to retries.resolve("ClientRateLimiter"),
         "ClientRateLimiterPartition" to retries.resolve("ClientRateLimiterPartition"),
-        "StaticPartitionMap" to smithyRuntimeCrate.resolve("static_partition_map::StaticPartitionMap"),
+        "StaticPartitionMap" to smithyRuntime.resolve("static_partition_map::StaticPartitionMap"),
     )
 
     override fun section(section: ServiceRuntimePluginSection): Writable = writable {
-        when (section) {
-            is ServiceRuntimePluginSection.DeclareSingletons -> {
-                // TODO(enableNewSmithyRuntimeCleanup) We can use the standard library's `OnceCell` once we upgrade the
-                //    MSRV to 1.70
-                rustTemplate(
-                    """
-                    static TOKEN_BUCKET: #{StaticPartitionMap}<#{TokenBucketPartition}, #{TokenBucket}> = #{StaticPartitionMap}::new();
-                    static CLIENT_RATE_LIMITER: #{StaticPartitionMap}<#{ClientRateLimiterPartition}, #{ClientRateLimiter}> = #{StaticPartitionMap}::new();
-                    """,
-                    *codegenScope,
-                )
+        if (runtimeMode.generateOrchestrator) {
+            when (section) {
+                is ServiceRuntimePluginSection.DeclareSingletons -> {
+                    // TODO(enableNewSmithyRuntimeCleanup) We can use the standard library's `OnceCell` once we upgrade the
+                    //    MSRV to 1.70
+                    rustTemplate(
+                        """
+                        static TOKEN_BUCKET: #{StaticPartitionMap}<#{TokenBucketPartition}, #{TokenBucket}> = #{StaticPartitionMap}::new();
+                        static CLIENT_RATE_LIMITER: #{StaticPartitionMap}<#{ClientRateLimiterPartition}, #{ClientRateLimiter}> = #{StaticPartitionMap}::new();
+                        """,
+                        *codegenScope,
+                    )
+                }
+
+                else -> emptySection
             }
-            else -> emptySection
         }
     }
 }

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt

@@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.client.smithy.customize
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule
+import software.amazon.smithy.rust.codegen.client.smithy.customizations.ConnectionPoisoningRuntimePluginCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.EndpointPrefixGenerator
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.HttpChecksumRequiredGenerator
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.HttpVersionListCustomization
@@ -82,7 +83,7 @@ class RequiredCustomizations : ClientCodegenDecorator {
         rustCrate.mergeFeature(TestUtilFeature)
 
         // Re-export resiliency types
-        ResiliencyReExportCustomization(codegenContext.runtimeConfig).extras(rustCrate)
+        ResiliencyReExportCustomization(codegenContext).extras(rustCrate)
 
         rustCrate.withModule(ClientRustModule.Primitives) {
             pubUseSmithyPrimitives(codegenContext, codegenContext.model)(this)
@@ -102,7 +103,9 @@ class RequiredCustomizations : ClientCodegenDecorator {
         codegenContext: ClientCodegenContext,
         baseCustomizations: List<ServiceRuntimePluginCustomization>,
     ): List<ServiceRuntimePluginCustomization> = if (codegenContext.smithyRuntimeMode.generateOrchestrator) {
-        baseCustomizations + ResiliencyServiceRuntimePluginCustomization(codegenContext)
+        baseCustomizations +
+            ResiliencyServiceRuntimePluginCustomization(codegenContext) +
+            ConnectionPoisoningRuntimePluginCustomization(codegenContext)
     } else {
         baseCustomizations
     }

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

@@ -67,13 +67,12 @@ sealed class ServiceRuntimePluginSection(name: String) : Section(name) {
     data class RegisterInterceptor(val interceptorRegistrarName: String) : ServiceRuntimePluginSection("RegisterInterceptor") {
         /** Generates the code to register an interceptor */
         fun registerInterceptor(runtimeConfig: RuntimeConfig, writer: RustWriter, interceptor: Writable) {
-            val smithyRuntimeApi = RuntimeType.smithyRuntimeApi(runtimeConfig)
             writer.rustTemplate(
                 """
                 $interceptorRegistrarName.register(#{SharedInterceptor}::new(#{interceptor}) as _);
                 """,
                 "interceptor" to interceptor,
-                "SharedInterceptor" to smithyRuntimeApi.resolve("client::interceptors::SharedInterceptor"),
+                "SharedInterceptor" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::interceptors::SharedInterceptor"),
             )
         }
     }

codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ResiliencyConfigCustomizationTest.kt

@@ -40,7 +40,7 @@ internal class ResiliencyConfigCustomizationTest {
         val codegenContext = testClientCodegenContext(model, settings = project.clientRustSettings())
 
         stubConfigProject(codegenContext, ResiliencyConfigCustomization(codegenContext), project)
-        ResiliencyReExportCustomization(codegenContext.runtimeConfig).extras(project)
+        ResiliencyReExportCustomization(codegenContext).extras(project)
         project.compileAndTest()
     }
 }