feature: Invocation ID interceptor (#2626)

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here -->
part of #1793 

## Description
<!--- Describe your changes in detail -->
This adds an interceptor for AWS SDK requests. The interceptor is run
just before the retry loop and adds a header with name
`amz-sdk-invocation-id` and value that's a UUID. AWS services use this
identifier to more efficiently process requests.

## 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. -->
This change includes tests

----

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

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

@@ -18,6 +18,7 @@ aws-types = { path = "../aws-types" }
 http = "0.2.3"
 percent-encoding = "2.1.0"
 tracing = "0.1"
+uuid = { version = "1", features = ["v4", "fast-rng"] }
 
 [dev-dependencies]
 aws-smithy-protocol-test = { path = "../../../rust-runtime/aws-smithy-protocol-test" }

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

@@ -0,0 +1,91 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_smithy_runtime_api::client::interceptors::error::BoxError;
+use aws_smithy_runtime_api::client::interceptors::{Interceptor, InterceptorContext};
+use aws_smithy_runtime_api::client::orchestrator::{HttpRequest, HttpResponse};
+use aws_smithy_runtime_api::config_bag::ConfigBag;
+use http::{HeaderName, HeaderValue};
+use uuid::Uuid;
+
+#[allow(clippy::declare_interior_mutable_const)] // we will never mutate this
+const AMZ_SDK_INVOCATION_ID: HeaderName = HeaderName::from_static("amz-sdk-invocation-id");
+
+/// This interceptor generates a UUID and attaches it to all request attempts made as part of this operation.
+#[non_exhaustive]
+#[derive(Debug)]
+pub struct InvocationIdInterceptor {
+ id: HeaderValue,
+}
+
+impl InvocationIdInterceptor {
+ /// Creates a new `InvocationIdInterceptor`
+ pub fn new() -> Self {
+ Self::default()
+ }
+}
+
+impl Default for InvocationIdInterceptor {
+ fn default() -> Self {
+ let id = Uuid::new_v4();
+ let id = id
+ .to_string()
+ .parse()
+ .expect("UUIDs always produce a valid header value");
+ Self { id }
+ }
+}
+
+impl Interceptor<HttpRequest, HttpResponse> for InvocationIdInterceptor {
+ fn modify_before_retry_loop(
+ &self,
+ context: &mut InterceptorContext<HttpRequest, HttpResponse>,
+ _cfg: &mut ConfigBag,
+ ) -> Result<(), BoxError> {
+ let headers = context.request_mut()?.headers_mut();
+ headers.append(AMZ_SDK_INVOCATION_ID, self.id.clone());
+ Ok(())
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use crate::invocation_id::InvocationIdInterceptor;
+ use aws_smithy_http::body::SdkBody;
+ use aws_smithy_runtime_api::client::interceptors::{Interceptor, InterceptorContext};
+ use aws_smithy_runtime_api::client::orchestrator::{HttpRequest, HttpResponse};
+ use aws_smithy_runtime_api::config_bag::ConfigBag;
+ use aws_smithy_runtime_api::type_erasure::TypedBox;
+ use http::HeaderValue;
+
+ fn expect_header<'a>(
+ context: &'a InterceptorContext<HttpRequest, HttpResponse>,
+ header_name: &str,
+ ) -> &'a HeaderValue {
+ context
+ .request()
+ .unwrap()
+ .headers()
+ .get(header_name)
+ .unwrap()
+ }
+
+ #[test]
+ fn test_id_is_generated_and_set() {
+ let mut context = InterceptorContext::new(TypedBox::new("doesntmatter").erase());
+ context.set_request(http::Request::builder().body(SdkBody::empty()).unwrap());
+
+ let mut config = ConfigBag::base();
+ let interceptor = InvocationIdInterceptor::new();
+ interceptor
+ .modify_before_retry_loop(&mut context, &mut config)
+ .unwrap();
+
+ let header = expect_header(&context, "amz-sdk-invocation-id");
+ assert_eq!(&interceptor.id, header);
+ // UUID should include 32 chars and 4 dashes
+ assert_eq!(interceptor.id.len(), 36);
+ }
+}

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

@@ -24,3 +24,6 @@ pub mod recursion_detection;
 
 /// Supporting code for user agent headers in the AWS SDK.
 pub mod user_agent;
+
+/// Supporting code for invocation ID headers in the AWS SDK.
+pub mod invocation_id;

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

@@ -53,6 +53,7 @@ val DECORATORS: List<ClientCodegenDecorator> = listOf(
  AwsRequestIdDecorator(),
  DisabledAuthDecorator(),
  RecursionDetectionDecorator(),
+ InvocationIdDecorator(),
  ),
 
  // Service specific decorators

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

@@ -0,0 +1,43 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rustsdk
+
+import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
+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.util.letIf
+
+class InvocationIdDecorator : ClientCodegenDecorator {
+ override val name: String get() = "InvocationIdDecorator"
+ override val order: Byte get() = 0
+ override fun serviceRuntimePluginCustomizations(
+ codegenContext: ClientCodegenContext,
+ baseCustomizations: List<ServiceRuntimePluginCustomization>,
+ ): List<ServiceRuntimePluginCustomization> =
+ baseCustomizations.letIf(codegenContext.settings.codegenConfig.enableNewSmithyRuntime) {
+ it + listOf(InvocationIdRuntimePluginCustomization(codegenContext))
+ }
+}
+
+private class InvocationIdRuntimePluginCustomization(
+ private val codegenContext: ClientCodegenContext,
+) : ServiceRuntimePluginCustomization() {
+ override fun section(section: ServiceRuntimePluginSection): Writable = writable {
+ if (section is ServiceRuntimePluginSection.AdditionalConfig) {
+ section.registerInterceptor(codegenContext.runtimeConfig, this) {
+ rust(
+ "#T::new()",
+ AwsRuntimeType.awsRuntime(codegenContext.runtimeConfig)
+ .resolve("invocation_id::InvocationIdInterceptor"),
+ )
+ }
+ }
+ }
+}