Unverified Commit f271da4c authored by Zelda Hessler's avatar Zelda Hessler Committed by GitHub
Browse files

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: default avatarJohn DiSanti <jdisanti@amazon.com>
parent e07aa766

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

+1 −0
Original line number Diff line number Diff line
@@ -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 → 100644
+91 −0
Original line number Diff line number Diff line 
/*

 * 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

+3 −0
Original line number Diff line number Diff line
@@ -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

+1 −0
Original line number Diff line number Diff line
@@ -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 → 100644
+43 −0
Original line number Diff line number Diff line 
/*

 * 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"),

                )

            }

        }

    }

}