Skip to content

Commit

Permalink
Allow canary-runner to specify lambda's memory size
Browse files Browse the repository at this point in the history
  • Loading branch information
@ysaito1001
ysaito1001 committed Mar 8, 2024
1 parent 75ed6fd commit eebe8af
Showing 1 changed file with 42 additions and 24 deletions.
66 changes: 42 additions & 24 deletions tools/ci-cdk/canary-runner/src/run.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -89,6 +91,10 @@ pub struct RunArgs {
#[clap(long)]
expected_speech_text_by_transcribe: Option<String>,

/// Memory allocated for the function.
#[clap(long)]
lambda_function_memory_size_in_mb: Option<i32>,

/// File path to a CDK outputs JSON file. This can be used instead
/// of all the --lambda... args.
#[clap(long)]
Expand Down Expand Up @@ -122,6 +128,7 @@ struct Options {
sdk_path: Option<PathBuf>,
musl: bool,
expected_speech_text_by_transcribe: Option<String>,
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,
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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!())?;
Expand Down Expand Up @@ -356,37 +364,43 @@ 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,
),
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
Expand All @@ -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"))?;
Expand Down Expand Up @@ -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]
Expand All @@ -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,
Expand All @@ -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",
])
Expand Down Expand Up @@ -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(),
Expand Down

0 comments on commit eebe8af

Please sign in to comment.