Skip to content

Commit

Permalink
Add feature ID for S3 Transfer Manager (#3921)
Browse files Browse the repository at this point in the history
## Motivation and Context
This PR introduces the `AwsSdkFeature` enum, which will be used for AWS
SDK-specific feature identifiers. For now, it includes a single ID for
the S3 Transfer Manager. This approach is similar to the existing
[SmithySdkFeature](https://github.com/smithy-lang/smithy-rs/blob/aac9becfd469e2479f68c61fd4c8074ddf755482/rust-runtime/aws-smithy-runtime/src/client/sdk_feature.rs#L10)
in terms of its module and crate structure (that is for generic client
but `AwsSdkFeature` is AWS specific); the `sdk_feature` module is hidden
as it is intended for internal use only.

## Testing
Added an integration test to verify tracking a business metric for
Transfer Manager. Since Transfer Manager is a high-level library, we do
not track its business metric directly within the `smithy-rs`
repository. Instead, [an s3
client](https://github.com/awslabs/aws-s3-transfer-manager-rs/blob/main/aws-s3-transfer-manager/src/config.rs#L21)
used in
[aws-s3-transfer-manager-rs](https://github.com/awslabs/aws-s3-transfer-manager-rs)
can configure the s3 client as demonstrated in the integration test to
track the metric.

----

_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: Zelda Hessler <zhessler@amazon.com>
  • Loading branch information
ysaito1001 and Velfi authored Nov 26, 2024
1 parent aac9bec commit e54cc56
Show file tree
Hide file tree
Showing 10 changed files with 131 additions and 26 deletions.
4 changes: 2 additions & 2 deletions aws/rust-runtime/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion aws/rust-runtime/aws-runtime/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "aws-runtime"
version = "1.4.3"
version = "1.4.4"
authors = ["AWS Rust SDK Team <aws-sdk-rust@amazon.com>"]
description = "Runtime support code for the AWS SDK. This crate isn't intended to be used directly."
edition = "2021"
Expand Down
4 changes: 4 additions & 0 deletions aws/rust-runtime/aws-runtime/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,10 @@ pub mod invocation_id;
/// Supporting code for request metadata headers in the AWS SDK.
pub mod request_info;

/// AWS SDK feature identifies.
#[doc(hidden)]
pub mod sdk_feature;

/// Interceptor that determines the clock skew between the client and service.
pub mod service_clock_skew;

Expand Down
18 changes: 18 additions & 0 deletions aws/rust-runtime/aws-runtime/src/sdk_feature.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

use aws_smithy_types::config_bag::{Storable, StoreAppend};

/// IDs for the features that may be used in the AWS SDK
#[non_exhaustive]
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum AwsSdkFeature {
/// Indicates that an operation was called by the S3 Transfer Manager
S3Transfer,
}

impl Storable for AwsSdkFeature {
type Storer = StoreAppend<Self>;
}
8 changes: 8 additions & 0 deletions aws/rust-runtime/aws-runtime/src/user_agent/interceptor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use aws_smithy_types::config_bag::ConfigBag;
use aws_types::app_name::AppName;
use aws_types::os_shim_internal::Env;

use crate::sdk_feature::AwsSdkFeature;
use crate::user_agent::metrics::ProvideBusinessMetric;
use crate::user_agent::{AdditionalMetadata, ApiMetadata, AwsUserAgent, InvalidMetadataValue};

Expand Down Expand Up @@ -138,6 +139,13 @@ impl Intercept for UserAgentInterceptor {
.map(|m| ua.add_business_metric(m));
}

let aws_sdk_features = cfg.load::<AwsSdkFeature>();
for aws_sdk_feature in aws_sdk_features {
aws_sdk_feature
.provide_business_metric()
.map(|m| ua.add_business_metric(m));
}

let maybe_connector_metadata = runtime_components
.http_client()
.and_then(|c| c.connector_metadata());
Expand Down
10 changes: 10 additions & 0 deletions aws/rust-runtime/aws-runtime/src/user_agent/metrics.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
* SPDX-License-Identifier: Apache-2.0
*/

use crate::sdk_feature::AwsSdkFeature;
use aws_smithy_runtime::client::sdk_feature::SmithySdkFeature;
use once_cell::sync::Lazy;
use std::borrow::Cow;
Expand Down Expand Up @@ -155,6 +156,15 @@ impl ProvideBusinessMetric for SmithySdkFeature {
}
}

impl ProvideBusinessMetric for AwsSdkFeature {
fn provide_business_metric(&self) -> Option<BusinessMetric> {
use AwsSdkFeature::*;
match self {
S3Transfer => Some(BusinessMetric::S3Transfer),
}
}
}

#[derive(Clone, Debug, Default)]
pub(super) struct BusinessMetrics(Vec<BusinessMetric>);

Expand Down
63 changes: 63 additions & 0 deletions aws/sdk/integration-tests/s3/tests/business_metrics.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

use aws_config::Region;
use aws_runtime::{
sdk_feature::AwsSdkFeature, user_agent::test_util::assert_ua_contains_metric_values,
};
use aws_sdk_s3::{
config::{Intercept, IntoShared},
primitives::ByteStream,
Client, Config,
};
use aws_smithy_runtime::client::http::test_util::capture_request;

#[derive(Debug)]
struct TransferManagerFeatureInterceptor;

impl Intercept for TransferManagerFeatureInterceptor {
fn name(&self) -> &'static str {
"TransferManagerFeature"
}

fn read_before_execution(
&self,
_ctx: &aws_sdk_s3::config::interceptors::BeforeSerializationInterceptorContextRef<'_>,
cfg: &mut aws_sdk_s3::config::ConfigBag,
) -> Result<(), aws_sdk_s3::error::BoxError> {
cfg.interceptor_state()
.store_append(AwsSdkFeature::S3Transfer);
Ok(())
}
}

#[tokio::test]
async fn test_track_metric_for_s3_transfer_manager() {
let (http_client, captured_request) = capture_request(None);
let mut conf_builder = Config::builder()
.region(Region::new("us-east-1"))
.http_client(http_client.clone())
.with_test_defaults();
// The S3 Transfer Manager uses a passed-in S3 client SDK for operations.
// By configuring an interceptor at the client level to track metrics,
// all operations executed by the client will automatically include the metric.
// This eliminates the need to apply `.config_override` on individual operations
// to insert the `TransferManagerFeatureInterceptor`.
conf_builder.push_interceptor(TransferManagerFeatureInterceptor.into_shared());
let client = Client::from_conf(conf_builder.build());

let _ = client
.put_object()
.bucket("doesnotmatter")
.key("doesnotmatter")
.body(ByteStream::from_static("Hello, world".as_bytes()))
.send()
.await
.unwrap();

let expected_req = captured_request.expect_request();
let user_agent = expected_req.headers().get("x-amz-user-agent").unwrap();
assert_ua_contains_metric_values(user_agent, &["G"]);
}
44 changes: 22 additions & 22 deletions rust-runtime/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion rust-runtime/aws-smithy-runtime/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "aws-smithy-runtime"
version = "1.7.3"
version = "1.7.4"
authors = ["AWS Rust SDK Team <aws-sdk-rust@amazon.com>", "Zelda Hessler <zhessler@amazon.com>"]
description = "The new smithy runtime crate"
edition = "2021"
Expand Down
2 changes: 2 additions & 0 deletions rust-runtime/aws-smithy-runtime/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,9 @@ pub mod interceptors;
/// Stalled stream protection for clients
pub mod stalled_stream_protection;

/// Generic Smithy SDK feature identifies.
#[doc(hidden)]
pub mod sdk_feature;

/// Smithy support-code for code generated waiters.
pub mod waiters;

0 comments on commit e54cc56

Please sign in to comment.