From eebe8afb6b50449fdc9c8702bef9a66fd9787d8a Mon Sep 17 00:00:00 2001 From: ysaito1001 Date: Fri, 8 Mar 2024 09:08:35 -0600 Subject: [PATCH] Allow `canary-runner` to specify lambda's memory size --- tools/ci-cdk/canary-runner/src/run.rs | 66 +++++++++++++++++---------- 1 file changed, 42 insertions(+), 24 deletions(-) diff --git a/tools/ci-cdk/canary-runner/src/run.rs b/tools/ci-cdk/canary-runner/src/run.rs index b4394c6348..85fb004816 100644 --- a/tools/ci-cdk/canary-runner/src/run.rs +++ b/tools/ci-cdk/canary-runner/src/run.rs @@ -36,6 +36,8 @@ use aws_sdk_lambda as lambda; use aws_sdk_s3 as s3; use std::collections::HashMap; +const DEFAULT_LAMBDA_FUNCTION_MEMORY_SIZE_IN_MB: i32 = 512; + lazy_static::lazy_static! { // Occasionally, a breaking change introduced in smithy-rs will cause the canary to fail // for older versions of the SDK since the canary is in the smithy-rs repository and will @@ -89,6 +91,10 @@ pub struct RunArgs { #[clap(long)] expected_speech_text_by_transcribe: Option, + /// Memory allocated for the function. + #[clap(long)] + lambda_function_memory_size_in_mb: Option, + /// File path to a CDK outputs JSON file. This can be used instead /// of all the --lambda... args. #[clap(long)] @@ -122,6 +128,7 @@ struct Options { sdk_path: Option, musl: bool, expected_speech_text_by_transcribe: Option, + lambda_function_memory_size_in_mb: i32, lambda_code_s3_bucket_name: String, lambda_test_s3_bucket_name: String, lambda_test_s3_mrap_bucket_arn: String, @@ -161,6 +168,9 @@ impl Options { sdk_path: run_opt.sdk_path, musl: run_opt.musl, expected_speech_text_by_transcribe: run_opt.expected_speech_text_by_transcribe, + lambda_function_memory_size_in_mb: run_opt + .lambda_function_memory_size_in_mb + .unwrap_or(DEFAULT_LAMBDA_FUNCTION_MEMORY_SIZE_IN_MB), lambda_code_s3_bucket_name: value.inner.lambda_code_s3_bucket_name, lambda_test_s3_bucket_name: value.inner.lambda_test_s3_bucket_name, lambda_test_s3_mrap_bucket_arn: value.inner.lambda_test_s3_mrap_bucket_arn, @@ -174,6 +184,9 @@ impl Options { sdk_path: run_opt.sdk_path, musl: run_opt.musl, expected_speech_text_by_transcribe: run_opt.expected_speech_text_by_transcribe, + lambda_function_memory_size_in_mb: run_opt + .lambda_function_memory_size_in_mb + .unwrap_or(DEFAULT_LAMBDA_FUNCTION_MEMORY_SIZE_IN_MB), lambda_code_s3_bucket_name: run_opt.lambda_code_s3_bucket_name.expect("required"), lambda_test_s3_bucket_name: run_opt.lambda_test_s3_bucket_name.expect("required"), lambda_test_s3_mrap_bucket_arn: run_opt @@ -283,12 +296,7 @@ async fn run_canary(options: &Options, config: &aws_config::SdkConfig) -> Result lambda_client.clone(), bundle_name, bundle_file_name, - &options.lambda_execution_role_arn, - options.expected_speech_text_by_transcribe.as_ref(), - &options.lambda_code_s3_bucket_name, - &options.lambda_test_s3_bucket_name, - &options.lambda_test_s3_mrap_bucket_arn, - &options.lambda_test_s3_express_bucket_name, + options, ) .await .context(here!())?; @@ -356,27 +364,27 @@ async fn upload_bundle( Ok(()) } -#[allow(clippy::too_many_arguments)] async fn create_lambda_fn( lambda_client: lambda::Client, bundle_name: &str, bundle_file_name: &str, - execution_role: &str, - expected_speech_text_by_transcribe: Option<&String>, - code_s3_bucket: &str, - test_s3_bucket: &str, - test_s3_mrap_bucket_arn: &str, - test_s3_express_bucket: &str, + options: &Options, ) -> Result<()> { use lambda::types::*; - let mut env_builder = match expected_speech_text_by_transcribe { + let mut env_builder = match &options.expected_speech_text_by_transcribe { Some(expected_speech_text_by_transcribe) => Environment::builder() .variables("RUST_BACKTRACE", "1") .variables("RUST_LOG", "info") - .variables("CANARY_S3_BUCKET_NAME", test_s3_bucket) - .variables("CANARY_S3_MRAP_BUCKET_ARN", test_s3_mrap_bucket_arn) - .variables("CANARY_S3_EXPRESS_BUCKET_NAME", test_s3_express_bucket) + .variables("CANARY_S3_BUCKET_NAME", &options.lambda_test_s3_bucket_name) + .variables( + "CANARY_S3_MRAP_BUCKET_ARN", + &options.lambda_test_s3_mrap_bucket_arn, + ) + .variables( + "CANARY_S3_EXPRESS_BUCKET_NAME", + &options.lambda_test_s3_express_bucket_name, + ) .variables( "CANARY_EXPECTED_TRANSCRIBE_RESULT", expected_speech_text_by_transcribe, @@ -384,9 +392,15 @@ async fn create_lambda_fn( None => Environment::builder() .variables("RUST_BACKTRACE", "1") .variables("RUST_LOG", "info") - .variables("CANARY_S3_BUCKET_NAME", test_s3_bucket) - .variables("CANARY_S3_MRAP_BUCKET_ARN", test_s3_mrap_bucket_arn) - .variables("CANARY_S3_EXPRESS_BUCKET_NAME", test_s3_express_bucket), + .variables("CANARY_S3_BUCKET_NAME", &options.lambda_test_s3_bucket_name) + .variables( + "CANARY_S3_MRAP_BUCKET_ARN", + &options.lambda_test_s3_mrap_bucket_arn, + ) + .variables( + "CANARY_S3_EXPRESS_BUCKET_NAME", + &options.lambda_test_s3_express_bucket_name, + ), }; // TODO(Post S3Express release): Delete this once S3 Express has been released and its canary is on by default @@ -398,17 +412,18 @@ async fn create_lambda_fn( .create_function() .function_name(bundle_name) .runtime(Runtime::Providedal2) - .role(execution_role) + .role(&options.lambda_execution_role_arn) .handler("aws-sdk-rust-lambda-canary") .code( FunctionCode::builder() - .s3_bucket(code_s3_bucket) + .s3_bucket(&options.lambda_code_s3_bucket_name) .s3_key(bundle_file_name) .build(), ) .publish(true) .environment(env_builder.build()) .timeout(180) + .memory_size(options.lambda_function_memory_size_in_mb) .send() .await .context(here!("failed to create canary Lambda function"))?; @@ -507,8 +522,7 @@ async fn delete_lambda_fn(lambda_client: lambda::Client, bundle_name: &str) -> R #[cfg(test)] mod tests { - use crate::run::Options; - use crate::run::RunArgs; + use crate::run::{Options, RunArgs, DEFAULT_LAMBDA_FUNCTION_MEMORY_SIZE_IN_MB}; use clap::Parser; #[test] @@ -520,6 +534,7 @@ mod tests { sdk_path: Some("artifact-aws-sdk-rust/sdk".into()), musl: false, expected_speech_text_by_transcribe: Some("Good day to you transcribe.".to_owned()), + lambda_function_memory_size_in_mb: Some(1024), cdk_output: Some("../cdk-outputs.json".into()), lambda_code_s3_bucket_name: None, lambda_test_s3_bucket_name: None, @@ -533,6 +548,8 @@ mod tests { "artifact-aws-sdk-rust/sdk", "--expected-speech-text-by-transcribe", "Good day to you transcribe.", + "--lambda-function-memory-size-in-mb", + "1024", "--cdk-output", "../cdk-outputs.json", ]) @@ -567,6 +584,7 @@ mod tests { sdk_path: Some("artifact-aws-sdk-rust/sdk".into()), musl: false, expected_speech_text_by_transcribe: Some("Good day to you transcribe.".to_owned()), + lambda_function_memory_size_in_mb: DEFAULT_LAMBDA_FUNCTION_MEMORY_SIZE_IN_MB, lambda_code_s3_bucket_name: "bucket-for-code".to_owned(), lambda_test_s3_bucket_name: "bucket-for-test".to_owned(), lambda_execution_role_arn: "arn:aws:lambda::role/exe-role".to_owned(),