Skip to content

Commit f742d2f

Browse files
Zelda Hesslerjdisanti
authored andcommitted
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>
1 parent c02f854 commit f742d2f

File tree

5 files changed

+139
-0
lines changed

5 files changed

+139
-0
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ aws-types = { path = "../aws-types" }
1818
http = "0.2.3"
1919
percent-encoding = "2.1.0"
2020
tracing = "0.1"
21+
uuid = { version = "1", features = ["v4", "fast-rng"] }
2122

2223
[dev-dependencies]
2324
aws-smithy-protocol-test = { path = "../../../rust-runtime/aws-smithy-protocol-test" }
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
use aws_smithy_runtime_api::client::interceptors::error::BoxError;
7+
use aws_smithy_runtime_api::client::interceptors::{Interceptor, InterceptorContext};
8+
use aws_smithy_runtime_api::client::orchestrator::{HttpRequest, HttpResponse};
9+
use aws_smithy_runtime_api::config_bag::ConfigBag;
10+
use http::{HeaderName, HeaderValue};
11+
use uuid::Uuid;
12+
13+
#[allow(clippy::declare_interior_mutable_const)] // we will never mutate this
14+
const AMZ_SDK_INVOCATION_ID: HeaderName = HeaderName::from_static("amz-sdk-invocation-id");
15+
16+
/// This interceptor generates a UUID and attaches it to all request attempts made as part of this operation.
17+
#[non_exhaustive]
18+
#[derive(Debug)]
19+
pub struct InvocationIdInterceptor {
20+
id: HeaderValue,
21+
}
22+
23+
impl InvocationIdInterceptor {
24+
/// Creates a new `InvocationIdInterceptor`
25+
pub fn new() -> Self {
26+
Self::default()
27+
}
28+
}
29+
30+
impl Default for InvocationIdInterceptor {
31+
fn default() -> Self {
32+
let id = Uuid::new_v4();
33+
let id = id
34+
.to_string()
35+
.parse()
36+
.expect("UUIDs always produce a valid header value");
37+
Self { id }
38+
}
39+
}
40+
41+
impl Interceptor<HttpRequest, HttpResponse> for InvocationIdInterceptor {
42+
fn modify_before_retry_loop(
43+
&self,
44+
context: &mut InterceptorContext<HttpRequest, HttpResponse>,
45+
_cfg: &mut ConfigBag,
46+
) -> Result<(), BoxError> {
47+
let headers = context.request_mut()?.headers_mut();
48+
headers.append(AMZ_SDK_INVOCATION_ID, self.id.clone());
49+
Ok(())
50+
}
51+
}
52+
53+
#[cfg(test)]
54+
mod tests {
55+
use crate::invocation_id::InvocationIdInterceptor;
56+
use aws_smithy_http::body::SdkBody;
57+
use aws_smithy_runtime_api::client::interceptors::{Interceptor, InterceptorContext};
58+
use aws_smithy_runtime_api::client::orchestrator::{HttpRequest, HttpResponse};
59+
use aws_smithy_runtime_api::config_bag::ConfigBag;
60+
use aws_smithy_runtime_api::type_erasure::TypedBox;
61+
use http::HeaderValue;
62+
63+
fn expect_header<'a>(
64+
context: &'a InterceptorContext<HttpRequest, HttpResponse>,
65+
header_name: &str,
66+
) -> &'a HeaderValue {
67+
context
68+
.request()
69+
.unwrap()
70+
.headers()
71+
.get(header_name)
72+
.unwrap()
73+
}
74+
75+
#[test]
76+
fn test_id_is_generated_and_set() {
77+
let mut context = InterceptorContext::new(TypedBox::new("doesntmatter").erase());
78+
context.set_request(http::Request::builder().body(SdkBody::empty()).unwrap());
79+
80+
let mut config = ConfigBag::base();
81+
let interceptor = InvocationIdInterceptor::new();
82+
interceptor
83+
.modify_before_retry_loop(&mut context, &mut config)
84+
.unwrap();
85+
86+
let header = expect_header(&context, "amz-sdk-invocation-id");
87+
assert_eq!(&interceptor.id, header);
88+
// UUID should include 32 chars and 4 dashes
89+
assert_eq!(interceptor.id.len(), 36);
90+
}
91+
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,3 +24,6 @@ pub mod recursion_detection;
2424

2525
/// Supporting code for user agent headers in the AWS SDK.
2626
pub mod user_agent;
27+
28+
/// Supporting code for invocation ID headers in the AWS SDK.
29+
pub mod invocation_id;

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ val DECORATORS: List<ClientCodegenDecorator> = listOf(
5353
AwsRequestIdDecorator(),
5454
DisabledAuthDecorator(),
5555
RecursionDetectionDecorator(),
56+
InvocationIdDecorator(),
5657
),
5758

5859
// Service specific decorators
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.rustsdk
7+
8+
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
9+
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
10+
import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginCustomization
11+
import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginSection
12+
import software.amazon.smithy.rust.codegen.core.rustlang.Writable
13+
import software.amazon.smithy.rust.codegen.core.rustlang.rust
14+
import software.amazon.smithy.rust.codegen.core.rustlang.writable
15+
import software.amazon.smithy.rust.codegen.core.util.letIf
16+
17+
class InvocationIdDecorator : ClientCodegenDecorator {
18+
override val name: String get() = "InvocationIdDecorator"
19+
override val order: Byte get() = 0
20+
override fun serviceRuntimePluginCustomizations(
21+
codegenContext: ClientCodegenContext,
22+
baseCustomizations: List<ServiceRuntimePluginCustomization>,
23+
): List<ServiceRuntimePluginCustomization> =
24+
baseCustomizations.letIf(codegenContext.settings.codegenConfig.enableNewSmithyRuntime) {
25+
it + listOf(InvocationIdRuntimePluginCustomization(codegenContext))
26+
}
27+
}
28+
29+
private class InvocationIdRuntimePluginCustomization(
30+
private val codegenContext: ClientCodegenContext,
31+
) : ServiceRuntimePluginCustomization() {
32+
override fun section(section: ServiceRuntimePluginSection): Writable = writable {
33+
if (section is ServiceRuntimePluginSection.AdditionalConfig) {
34+
section.registerInterceptor(codegenContext.runtimeConfig, this) {
35+
rust(
36+
"#T::new()",
37+
AwsRuntimeType.awsRuntime(codegenContext.runtimeConfig)
38+
.resolve("invocation_id::InvocationIdInterceptor"),
39+
)
40+
}
41+
}
42+
}
43+
}

0 commit comments

Comments
 (0)