Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Simplify default config and add default async sleep #3071

Merged
merged 6 commits into from
Oct 17, 2023
Merged

Simplify default config and add default async sleep #3071

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9b9757c
Simplify default config and add default async sleep
jdisanti Oct 14, 2023
ea3ae96
Incorporate feedback
jdisanti Oct 16, 2023
3ddb51c
Merge remote-tracking branch 'origin/main' into jdisanti-generic-defa…
jdisanti Oct 16, 2023
1bebbe7
Fix clippy lints
jdisanti Oct 16, 2023
652f2ec
Fix tests
jdisanti Oct 16, 2023
f0bf2b7
Merge branch 'main' into jdisanti-generic-default-sleep
jdisanti Oct 17, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions aws/rust-runtime/aws-config/src/imds/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,13 +238,14 @@ impl ImdsCommonRuntimePlugin {
fn new(
config: &ProviderConfig,
endpoint_resolver: ImdsEndpointResolver,
retry_config: &RetryConfig,
retry_config: RetryConfig,
timeout_config: TimeoutConfig,
) -> Self {
let mut layer = Layer::new("ImdsCommonRuntimePlugin");
layer.store_put(AuthSchemeOptionResolverParams::new(()));
layer.store_put(EndpointResolverParams::new(()));
layer.store_put(SensitiveOutput);
layer.store_put(retry_config);
layer.store_put(timeout_config);
layer.store_put(user_agent());

Expand All @@ -255,7 +256,7 @@ impl ImdsCommonRuntimePlugin {
.with_endpoint_resolver(Some(endpoint_resolver))
.with_interceptor(UserAgentInterceptor::new())
.with_retry_classifier(SharedRetryClassifier::new(ImdsResponseRetryClassifier))
.with_retry_strategy(Some(StandardRetryStrategy::new(retry_config)))
.with_retry_strategy(Some(StandardRetryStrategy::new()))
.with_time_source(Some(config.time_source()))
.with_sleep_impl(config.sleep_impl()),
}
Expand Down Expand Up @@ -423,7 +424,7 @@ impl Builder {
let common_plugin = SharedRuntimePlugin::new(ImdsCommonRuntimePlugin::new(
&config,
endpoint_resolver,
&retry_config,
retry_config,
timeout_config,
));
let operation = Operation::builder()
Expand Down
88 changes: 1 addition & 87 deletions .../amazon/smithy/rust/codegen/client/smithy/customizations/ResiliencyConfigCustomization.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,20 +7,16 @@ 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.ClientRustModule
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.client.smithy.generators.config.ConfigCustomization
import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
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.RuntimeType
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
import software.amazon.smithy.rust.codegen.core.util.sdkId

class ResiliencyConfigCustomization(private val codegenContext: ClientCodegenContext) : ConfigCustomization() {
class ResiliencyConfigCustomization(codegenContext: ClientCodegenContext) : ConfigCustomization() {
private val runtimeConfig = codegenContext.runtimeConfig
private val retryConfig = RuntimeType.smithyTypes(runtimeConfig).resolve("retry")
private val sleepModule = RuntimeType.smithyAsync(runtimeConfig).resolve("rt::sleep")
Expand All @@ -44,8 +40,6 @@ class ResiliencyConfigCustomization(private val codegenContext: ClientCodegenCon
"StandardRetryStrategy" to retries.resolve("strategy::StandardRetryStrategy"),
"SystemTime" to RuntimeType.std.resolve("time::SystemTime"),
"TimeoutConfig" to timeoutModule.resolve("TimeoutConfig"),
"TokenBucket" to retries.resolve("TokenBucket"),
"TokenBucketPartition" to retries.resolve("TokenBucketPartition"),
)

override fun section(section: ServiceConfig) =
Expand Down Expand Up @@ -281,57 +275,6 @@ class ResiliencyConfigCustomization(private val codegenContext: ClientCodegenCon
)
}

is ServiceConfig.BuilderBuild -> {
rustTemplate(
"""
if layer.load::<#{RetryConfig}>().is_none() {
layer.store_put(#{RetryConfig}::disabled());
}
let retry_config = layer.load::<#{RetryConfig}>().expect("set to default above").clone();

if layer.load::<#{RetryPartition}>().is_none() {
layer.store_put(#{RetryPartition}::new("${codegenContext.serviceShape.sdkId()}"));
}
let retry_partition = layer.load::<#{RetryPartition}>().expect("set to default above").clone();

if retry_config.has_retry() {
#{debug}!("using retry strategy with partition '{}'", retry_partition);
}

if retry_config.mode() == #{RetryMode}::Adaptive {
if let #{Some}(time_source) = self.runtime_components.time_source() {
let seconds_since_unix_epoch = time_source
.now()
.duration_since(#{SystemTime}::UNIX_EPOCH)
.expect("the present takes place after the UNIX_EPOCH")
.as_secs_f64();
let client_rate_limiter_partition = #{ClientRateLimiterPartition}::new(retry_partition.clone());
let client_rate_limiter = CLIENT_RATE_LIMITER.get_or_init(client_rate_limiter_partition, || {
#{ClientRateLimiter}::new(seconds_since_unix_epoch)
});
layer.store_put(client_rate_limiter);
}
}

// The token bucket is used for both standard AND adaptive retries.
let token_bucket_partition = #{TokenBucketPartition}::new(retry_partition);
let token_bucket = TOKEN_BUCKET.get_or_init(token_bucket_partition, #{TokenBucket}::default);
layer.store_put(token_bucket);

// TODO(enableNewSmithyRuntimeCleanup): Should not need to provide a default once smithy-rs##2770
// is resolved
if layer.load::<#{TimeoutConfig}>().is_none() {
layer.store_put(#{TimeoutConfig}::disabled());
}

self.runtime_components.set_retry_strategy(#{Some}(
#{SharedRetryStrategy}::new(#{StandardRetryStrategy}::new(&retry_config)))
);
""",
*codegenScope,
)
}

else -> emptySection
}
}
Expand Down Expand Up @@ -366,32 +309,3 @@ class ResiliencyReExportCustomization(codegenContext: ClientCodegenContext) {
}
}
}

class ResiliencyServiceRuntimePluginCustomization(codegenContext: ClientCodegenContext) : ServiceRuntimePluginCustomization() {
private val runtimeConfig = codegenContext.runtimeConfig
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 smithyRuntime.resolve("static_partition_map::StaticPartitionMap"),
)

override fun section(section: ServiceRuntimePluginSection): Writable = writable {
when (section) {
is ServiceRuntimePluginSection.DeclareSingletons -> {
rustTemplate(
"""
static TOKEN_BUCKET: #{StaticPartitionMap}<#{TokenBucketPartition}, #{TokenBucket}> = #{StaticPartitionMap}::new();
static CLIENT_RATE_LIMITER: #{StaticPartitionMap}<#{ClientRateLimiterPartition}, #{ClientRateLimiter}> = #{StaticPartitionMap}::new();
""",
*codegenScope,
)
}

else -> emptySection
}
}
}
2 changes: 0 additions & 2 deletions ...lin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ import software.amazon.smithy.rust.codegen.client.smithy.customizations.Intercep
import software.amazon.smithy.rust.codegen.client.smithy.customizations.MetadataCustomization
import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyConfigCustomization
import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyReExportCustomization
import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyServiceRuntimePluginCustomization
import software.amazon.smithy.rust.codegen.client.smithy.customizations.RetryClassifierConfigCustomization
import software.amazon.smithy.rust.codegen.client.smithy.customizations.RetryClassifierOperationCustomization
import software.amazon.smithy.rust.codegen.client.smithy.customizations.RetryClassifierServiceRuntimePluginCustomization
Expand Down Expand Up @@ -115,7 +114,6 @@ class RequiredCustomizations : ClientCodegenDecorator {
codegenContext: ClientCodegenContext,
baseCustomizations: List<ServiceRuntimePluginCustomization>,
): List<ServiceRuntimePluginCustomization> = baseCustomizations +
ResiliencyServiceRuntimePluginCustomization(codegenContext) +
ConnectionPoisoningRuntimePluginCustomization(codegenContext) +
RetryClassifierServiceRuntimePluginCustomization(codegenContext)
}
39 changes: 28 additions & 11 deletions ...tware/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
import software.amazon.smithy.rust.codegen.core.rustlang.stripOuter
import software.amazon.smithy.rust.codegen.core.rustlang.withBlockTemplate
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
Expand All @@ -50,9 +49,11 @@ import software.amazon.smithy.rust.codegen.core.smithy.expectRustMetadata
import software.amazon.smithy.rust.codegen.core.smithy.generators.getterName
import software.amazon.smithy.rust.codegen.core.smithy.generators.setterName
import software.amazon.smithy.rust.codegen.core.smithy.rustType
import software.amazon.smithy.rust.codegen.core.util.dq
import software.amazon.smithy.rust.codegen.core.util.inputShape
import software.amazon.smithy.rust.codegen.core.util.orNull
import software.amazon.smithy.rust.codegen.core.util.outputShape
import software.amazon.smithy.rust.codegen.core.util.sdkId
import software.amazon.smithy.rust.codegen.core.util.toSnakeCase

class FluentClientGenerator(
Expand Down Expand Up @@ -161,7 +162,7 @@ class FluentClientGenerator(
}
""",
*clientScope,
"base_client_runtime_plugins" to baseClientRuntimePluginsFn(runtimeConfig),
"base_client_runtime_plugins" to baseClientRuntimePluginsFn(codegenContext),
)
}

Expand Down Expand Up @@ -446,22 +447,36 @@ class FluentClientGenerator(
}
}

private fun baseClientRuntimePluginsFn(runtimeConfig: RuntimeConfig): RuntimeType =
private fun baseClientRuntimePluginsFn(codegenContext: ClientCodegenContext): RuntimeType = codegenContext.runtimeConfig.let { rc ->
RuntimeType.forInlineFun("base_client_runtime_plugins", ClientRustModule.config) {
val api = RuntimeType.smithyRuntimeApi(rc)
val rt = RuntimeType.smithyRuntime(rc)
rustTemplate(
"""
pub(crate) fn base_client_runtime_plugins(
mut config: crate::Config,
) -> #{RuntimePlugins} {
let mut configured_plugins = #{Vec}::new();
::std::mem::swap(&mut config.runtime_plugins, &mut configured_plugins);

let defaults = [
#{default_http_client_plugin}(),
#{default_retry_config_plugin}(${codegenContext.serviceShape.sdkId().dq()}),
#{default_sleep_impl_plugin}(),
#{default_time_source_plugin}(),
#{default_timeout_config_plugin}(),
].into_iter().flatten();

let mut plugins = #{RuntimePlugins}::new()
.with_client_plugin(#{default_http_client_plugin}())
// defaults
.with_client_plugins(defaults)
// user config
.with_client_plugin(
#{StaticRuntimePlugin}::new()
.with_config(config.config.clone())
.with_runtime_components(config.runtime_components.clone())
)
// codegen config
.with_client_plugin(crate::config::ServiceRuntimePlugin::new(config))
.with_client_plugin(#{NoAuthRuntimePlugin}::new());
for plugin in configured_plugins {
Expand All @@ -471,15 +486,17 @@ private fun baseClientRuntimePluginsFn(runtimeConfig: RuntimeConfig): RuntimeTyp
}
""",
*preludeScope,
"RuntimePlugins" to RuntimeType.runtimePlugins(runtimeConfig),
"NoAuthRuntimePlugin" to RuntimeType.smithyRuntime(runtimeConfig)
.resolve("client::auth::no_auth::NoAuthRuntimePlugin"),
"StaticRuntimePlugin" to RuntimeType.smithyRuntimeApi(runtimeConfig)
.resolve("client::runtime_plugin::StaticRuntimePlugin"),
"default_http_client_plugin" to RuntimeType.smithyRuntime(runtimeConfig)
.resolve("client::http::default_http_client_plugin"),
"default_http_client_plugin" to rt.resolve("client::defaults::default_http_client_plugin"),
"default_retry_config_plugin" to rt.resolve("client::defaults::default_retry_config_plugin"),
"default_sleep_impl_plugin" to rt.resolve("client::defaults::default_sleep_impl_plugin"),
"default_timeout_config_plugin" to rt.resolve("client::defaults::default_timeout_config_plugin"),
"default_time_source_plugin" to rt.resolve("client::defaults::default_time_source_plugin"),
"NoAuthRuntimePlugin" to rt.resolve("client::auth::no_auth::NoAuthRuntimePlugin"),
"RuntimePlugins" to RuntimeType.runtimePlugins(rc),
"StaticRuntimePlugin" to api.resolve("client::runtime_plugin::StaticRuntimePlugin"),
)
}
}

/**
* For a given `operation` shape, return a list of strings where each string describes the name and input type of one of
Expand Down
2 changes: 1 addition & 1 deletion ...zon/smithy/rust/codegen/client/smithy/customizations/ResiliencyConfigCustomizationTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ internal class ResiliencyConfigCustomizationTest {
project.withModule(ClientRustModule.config) {
ServiceRuntimePluginGenerator(codegenContext).render(
this,
listOf(ResiliencyServiceRuntimePluginCustomization(codegenContext)),
emptyList(),
)
}
ResiliencyReExportCustomization(codegenContext).extras(project)
Expand Down
8 changes: 6 additions & 2 deletions ...e/amazon/smithy/rust/codegen/client/smithy/customizations/SensitiveOutputDecoratorTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.client.smithy.customizations
import org.junit.jupiter.api.Test
import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
Expand All @@ -16,6 +17,8 @@ import software.amazon.smithy.rust.codegen.core.testutil.integrationTest

class SensitiveOutputDecoratorTest {
private fun codegenScope(runtimeConfig: RuntimeConfig): Array<Pair<String, Any>> = arrayOf(
"capture_test_logs" to CargoDependency.smithyRuntimeTestUtil(runtimeConfig).toType()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for starting to use this, which I did not when adding this test.

.resolve("test_util::capture_test_logs::capture_test_logs"),
"capture_request" to RuntimeType.captureRequest(runtimeConfig),
"SdkBody" to RuntimeType.sdkBody(runtimeConfig),
)
Expand Down Expand Up @@ -48,10 +51,10 @@ class SensitiveOutputDecoratorTest {
rustCrate.integrationTest("redacting_sensitive_response_body") {
val moduleName = codegenContext.moduleUseName()
Attribute.TokioTest.render(this)
Attribute.TracedTest.render(this)
rustTemplate(
"""
async fn redacting_sensitive_response_body() {
let (_logs, logs_rx) = #{capture_test_logs}();
let (http_client, _r) = #{capture_request}(Some(
http::Response::builder()
.status(200)
Expand All @@ -69,7 +72,8 @@ class SensitiveOutputDecoratorTest {
.await
.expect("success");

assert!(logs_contain("** REDACTED **"));
let log_contents = logs_rx.contents();
assert!(log_contents.contains("** REDACTED **"));
}
""",
*codegenScope(codegenContext.runtimeConfig),
Expand Down
10 changes: 8 additions & 2 deletions .../smithy/rust/codegen/client/smithy/generators/ConfigOverrideRuntimePluginGeneratorTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {
}

@Test
fun `operation overrides retry strategy`() {
fun `operation overrides retry config`() {
clientIntegrationTest(model) { clientCodegenContext, rustCrate ->
val runtimeConfig = clientCodegenContext.runtimeConfig
val codegenScope = arrayOf(
Expand All @@ -176,11 +176,13 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {
.resolve("client::retries::RetryClassifiers"),
"RuntimeComponentsBuilder" to RuntimeType.runtimeComponentsBuilder(runtimeConfig),
"RuntimePlugin" to RuntimeType.runtimePlugin(runtimeConfig),
"StandardRetryStrategy" to RuntimeType.smithyRuntime(runtimeConfig)
.resolve("client::retries::strategy::StandardRetryStrategy"),
"ShouldAttempt" to RuntimeType.smithyRuntimeApi(runtimeConfig)
.resolve("client::retries::ShouldAttempt"),
)
rustCrate.testModule {
unitTest("test_operation_overrides_retry_strategy") {
unitTest("test_operation_overrides_retry_config") {
rustTemplate(
"""
use #{RuntimePlugin};
Expand All @@ -205,6 +207,8 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {

// Emulate the merging of runtime components from runtime plugins that the orchestrator does
let runtime_components = #{RuntimeComponentsBuilder}::for_tests()
// emulate the default retry config plugin by setting a retry strategy
.with_retry_strategy(#{Some}(#{StandardRetryStrategy}::new()))
.merge_from(&client_config.runtime_components)
.merge_from(&retry_classifiers_component)
.build()
Expand All @@ -231,6 +235,8 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {

// Emulate the merging of runtime components from runtime plugins that the orchestrator does
let runtime_components = #{RuntimeComponentsBuilder}::for_tests()
// emulate the default retry config plugin by setting a retry strategy
.with_retry_strategy(#{Some}(#{StandardRetryStrategy}::new()))
.merge_from(&client_config.runtime_components)
.merge_from(&retry_classifiers_component)
.merge_from(&config_override.runtime_components)
Expand Down
Loading
Loading