From 16a92c453b97cc81d9402dd62d0afa73b0f6ab3b Mon Sep 17 00:00:00 2001 From: Zelda Hessler Date: Tue, 17 Dec 2024 10:46:02 -0600 Subject: [PATCH] add retry mode business metric tracking --- .../smithy/rustsdk/UserAgentDecoratorTest.kt | 117 ++++++++++++++++++ ...eatureTrackerRuntimePluginCustomization.kt | 38 ++++++ .../customize/RequiredCustomizations.kt | 4 +- rust-runtime/Cargo.lock | 40 +++--- .../inlineable/src/sdk_feature_tracker.rs | 43 +------ 5 files changed, 184 insertions(+), 58 deletions(-) create mode 100644 codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/RetryModeFeatureTrackerRuntimePluginCustomization.kt diff --git a/aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/UserAgentDecoratorTest.kt b/aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/UserAgentDecoratorTest.kt index 1f200faed0..aa6c4f881c 100644 --- a/aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/UserAgentDecoratorTest.kt +++ b/aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/UserAgentDecoratorTest.kt @@ -212,6 +212,123 @@ class UserAgentDecoratorTest { } } + @Test + fun `it emits business metrics for retry modes`() { + val model = + """ + namespace test + + use aws.auth#sigv4 + use aws.api#service + use aws.protocols#restJson1 + use smithy.rules#endpointRuleSet + + @auth([sigv4]) + @sigv4(name: "dontcare") + @endpointRuleSet({ + "version": "1.0", + "rules": [{ "type": "endpoint", "conditions": [], "endpoint": { "url": "https://example.com" } }], + "parameters": {} + }) + @service(sdkId: "dontcare") + @restJson1 + service TestService { version: "2023-01-01", operations: [SomeOperation] } + + @http(uri: "/SomeOperation", method: "GET") + @optionalAuth + operation SomeOperation { + input: SomeInput, + output: SomeOutput + } + + @input + structure SomeInput {} + + @output + structure SomeOutput {} + """.asSmithyModel() + + awsSdkIntegrationTest(model) { ctx, rustCrate -> + rustCrate.integrationTest("retry_mode_feature_tracker") { + val rc = ctx.runtimeConfig + val moduleName = ctx.moduleUseName() + + rust( + """ + use $moduleName::config::{Region, retry::RetryConfig}; + use $moduleName::{Client, Config}; + """, + ) + + tokioTest("should_emit_metric_in_user_agent_standard_mode") { + rustTemplate( + """ + let (http_client, rcvr) = #{capture_request}(#{None}); + let config = Config::builder() + .region(Region::new("us-east-1")) + .http_client(http_client.clone()) + .retry_config(RetryConfig::standard()) + .with_test_defaults() + .build(); + let client = Client::from_conf(config); + let _ = client.some_operation().send().await; + let expected_req = rcvr.expect_request(); + let user_agent = expected_req + .headers() + .get("x-amz-user-agent") + .unwrap(); + #{assert_ua_contains_metric_values}(user_agent, &["E"]); + """, + *preludeScope, + "assert_ua_contains_metric_values" to + AwsRuntimeType.awsRuntimeTestUtil(rc) + .resolve("user_agent::test_util::assert_ua_contains_metric_values"), + "capture_request" to RuntimeType.captureRequest(rc), + "disable_interceptor" to + RuntimeType.smithyRuntimeApiClient(rc) + .resolve("client::interceptors::disable_interceptor"), + "UserAgentInterceptor" to + AwsRuntimeType.awsRuntime(rc) + .resolve("user_agent::UserAgentInterceptor"), + ) + } + + tokioTest("should_emit_metric_in_user_agent_adaptive_mode") { + rustTemplate( + """ + let (http_client, rcvr) = #{capture_request}(#{None}); + let config = Config::builder() + .region(Region::new("us-east-1")) + .http_client(http_client.clone()) + .retry_config(RetryConfig::adaptive()) + .with_test_defaults() + .build(); + let client = Client::from_conf(config); + let _ = client.some_operation().send().await; + let expected_req = rcvr.expect_request(); + let user_agent = expected_req + .headers() + .get("x-amz-user-agent") + .unwrap(); + #{assert_ua_contains_metric_values}(user_agent, &["F"]); + """, + *preludeScope, + "assert_ua_contains_metric_values" to + AwsRuntimeType.awsRuntimeTestUtil(rc) + .resolve("user_agent::test_util::assert_ua_contains_metric_values"), + "capture_request" to RuntimeType.captureRequest(rc), + "disable_interceptor" to + RuntimeType.smithyRuntimeApiClient(rc) + .resolve("client::interceptors::disable_interceptor"), + "UserAgentInterceptor" to + AwsRuntimeType.awsRuntime(rc) + .resolve("user_agent::UserAgentInterceptor"), + ) + } + } + } + } + @Test fun `it emits business metric for checksum usage`() { val model = diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/RetryModeFeatureTrackerRuntimePluginCustomization.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/RetryModeFeatureTrackerRuntimePluginCustomization.kt new file mode 100644 index 0000000000..db02fba95d --- /dev/null +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/RetryModeFeatureTrackerRuntimePluginCustomization.kt @@ -0,0 +1,38 @@ +/* + * 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.InlineDependency +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 + +class RetryModeFeatureTrackerRuntimePluginCustomization(codegenContext: ClientCodegenContext) : + ServiceRuntimePluginCustomization() { + private val runtimeConfig = codegenContext.runtimeConfig + + override fun section(section: ServiceRuntimePluginSection): Writable = + writable { + when (section) { + is ServiceRuntimePluginSection.RegisterRuntimeComponents -> { + section.registerInterceptor(this) { + rust( + "#T::new()", + RuntimeType.forInlineDependency( + InlineDependency.sdkFeatureTracker(runtimeConfig), + ).resolve("retry_mode::RetryModeFeatureTrackerInterceptor"), + ) + } + } + + else -> emptySection + } + } +} diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt index a16f878002..3c11c75011 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt @@ -19,6 +19,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customizations.Resilien 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 +import software.amazon.smithy.rust.codegen.client.smithy.customizations.RetryModeFeatureTrackerRuntimePluginCustomization import software.amazon.smithy.rust.codegen.client.smithy.customizations.TimeSourceCustomization import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationCustomization import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginCustomization @@ -134,5 +135,6 @@ class RequiredCustomizations : ClientCodegenDecorator { ): List = baseCustomizations + ConnectionPoisoningRuntimePluginCustomization(codegenContext) + - RetryClassifierServiceRuntimePluginCustomization(codegenContext) + RetryClassifierServiceRuntimePluginCustomization(codegenContext) + + RetryModeFeatureTrackerRuntimePluginCustomization(codegenContext) } diff --git a/rust-runtime/Cargo.lock b/rust-runtime/Cargo.lock index f1931480c9..62a445b1a2 100644 --- a/rust-runtime/Cargo.lock +++ b/rust-runtime/Cargo.lock @@ -295,7 +295,7 @@ dependencies = [ name = "aws-smithy-cbor" version = "0.60.8" dependencies = [ - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "criterion", "minicbor", ] @@ -305,7 +305,7 @@ name = "aws-smithy-checksums" version = "0.60.13" dependencies = [ "aws-smithy-http 0.60.11", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "bytes", "bytes-utils", "crc32c", @@ -353,7 +353,7 @@ name = "aws-smithy-compression" version = "0.0.2" dependencies = [ "aws-smithy-runtime-api 1.7.3", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "bytes", "bytes-utils", "flate2", @@ -374,7 +374,7 @@ name = "aws-smithy-eventstream" version = "0.60.5" dependencies = [ "arbitrary", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "bytes", "bytes-utils", "crc32fast", @@ -397,9 +397,9 @@ name = "aws-smithy-experimental" version = "0.1.5" dependencies = [ "aws-smithy-async 1.2.1", - "aws-smithy-runtime 1.7.4", + "aws-smithy-runtime 1.7.6", "aws-smithy-runtime-api 1.7.3", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "h2 0.4.7", "http 1.2.0", "hyper 1.5.1", @@ -420,7 +420,7 @@ dependencies = [ "async-stream", "aws-smithy-eventstream 0.60.5", "aws-smithy-runtime-api 1.7.3", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "bytes", "bytes-utils", "futures-core", @@ -470,7 +470,7 @@ dependencies = [ "aws-smithy-http 0.60.11", "aws-smithy-json 0.61.1", "aws-smithy-runtime-api 1.7.3", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "aws-smithy-xml 0.60.9", "bytes", "futures-util", @@ -500,7 +500,7 @@ dependencies = [ "aws-smithy-http 0.60.11", "aws-smithy-http-server", "aws-smithy-json 0.61.1", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "aws-smithy-xml 0.60.9", "bytes", "futures", @@ -549,7 +549,7 @@ dependencies = [ name = "aws-smithy-json" version = "0.61.1" dependencies = [ - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "proptest", "serde_json", ] @@ -560,7 +560,7 @@ version = "0.2.1" dependencies = [ "aws-sdk-s3", "aws-smithy-runtime-api 1.7.3", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "tokio", ] @@ -604,7 +604,7 @@ dependencies = [ name = "aws-smithy-query" version = "0.60.7" dependencies = [ - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "urlencoding", ] @@ -642,14 +642,14 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.7.4" +version = "1.7.6" dependencies = [ "approx", "aws-smithy-async 1.2.1", "aws-smithy-http 0.60.11", "aws-smithy-protocol-test 0.63.0", "aws-smithy-runtime-api 1.7.3", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "bytes", "fastrand", "futures-util", @@ -680,7 +680,7 @@ name = "aws-smithy-runtime-api" version = "1.7.3" dependencies = [ "aws-smithy-async 1.2.1", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "bytes", "http 0.2.12", "http 1.2.0", @@ -736,7 +736,7 @@ dependencies = [ [[package]] name = "aws-smithy-types" -version = "1.2.10" +version = "1.2.11" dependencies = [ "base64 0.13.1", "base64-simd", @@ -773,7 +773,7 @@ name = "aws-smithy-types-convert" version = "0.60.8" dependencies = [ "aws-smithy-async 1.2.1", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "chrono", "futures-core", "time", @@ -785,7 +785,7 @@ version = "0.1.3" dependencies = [ "aws-smithy-http 0.60.11", "aws-smithy-runtime-api 1.7.3", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "bytes", "http 1.2.0", "tracing", @@ -1989,9 +1989,9 @@ dependencies = [ "aws-smithy-compression", "aws-smithy-http 0.60.11", "aws-smithy-json 0.61.1", - "aws-smithy-runtime 1.7.4", + "aws-smithy-runtime 1.7.6", "aws-smithy-runtime-api 1.7.3", - "aws-smithy-types 1.2.10", + "aws-smithy-types 1.2.11", "aws-smithy-xml 0.60.9", "bytes", "fastrand", diff --git a/rust-runtime/inlineable/src/sdk_feature_tracker.rs b/rust-runtime/inlineable/src/sdk_feature_tracker.rs index 710f4fcce5..acc653842b 100644 --- a/rust-runtime/inlineable/src/sdk_feature_tracker.rs +++ b/rust-runtime/inlineable/src/sdk_feature_tracker.rs @@ -177,18 +177,14 @@ pub(crate) mod waiter { pub(crate) mod retry_mode { use aws_smithy_runtime::client::sdk_feature::SmithySdkFeature; use aws_smithy_runtime_api::box_error::BoxError; - use aws_smithy_runtime_api::client::interceptors::context::BeforeSerializationInterceptorContextMut; - use aws_smithy_runtime_api::client::interceptors::{Intercept, SharedInterceptor}; - use aws_smithy_runtime_api::client::runtime_components::{ - RuntimeComponents, RuntimeComponentsBuilder, - }; - use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin; + use aws_smithy_runtime_api::client::interceptors::context::BeforeSerializationInterceptorContextRef; + use aws_smithy_runtime_api::client::interceptors::Intercept; + use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents; use aws_smithy_types::config_bag::ConfigBag; use aws_smithy_types::retry::{RetryConfig, RetryMode}; - use std::borrow::Cow; #[derive(Debug)] - struct RetryModeFeatureTrackerInterceptor; + pub(crate) struct RetryModeFeatureTrackerInterceptor; impl RetryModeFeatureTrackerInterceptor { pub(crate) fn new() -> Self { @@ -201,9 +197,9 @@ pub(crate) mod retry_mode { "RetryModeFeatureTrackerInterceptor" } - fn modify_before_serialization( + fn read_before_serialization( &self, - _context: &mut BeforeSerializationInterceptorContextMut<'_>, + _context: &BeforeSerializationInterceptorContextRef<'_>, _runtime_components: &RuntimeComponents, cfg: &mut ConfigBag, ) -> Result<(), BoxError> { @@ -221,31 +217,4 @@ pub(crate) mod retry_mode { Ok(()) } } - - #[derive(Debug)] - pub(crate) struct RetryModeFeatureTrackerRuntimePlugin { - runtime_components: RuntimeComponentsBuilder, - } - - impl RetryModeFeatureTrackerRuntimePlugin { - pub(crate) fn new() -> Self { - Self { - runtime_components: RuntimeComponentsBuilder::new( - "RetryModeFeatureTrackerRuntimePlugin", - ) - .with_interceptor(SharedInterceptor::new( - RetryModeFeatureTrackerInterceptor::new(), - )), - } - } - } - - impl RuntimePlugin for RetryModeFeatureTrackerRuntimePlugin { - fn runtime_components( - &self, - _: &RuntimeComponentsBuilder, - ) -> Cow<'_, RuntimeComponentsBuilder> { - Cow::Borrowed(&self.runtime_components) - } - } }