From c0d4475a23f70fdf8f8a1edb4d8f4b661c32eb94 Mon Sep 17 00:00:00 2001 From: Tom Kirchner Date: Sun, 23 Aug 2020 16:10:12 -0700 Subject: [PATCH 1/2] pubsys: check for copied AMIs in parallel We were waiting on the get_ami_id calls serially which led to an unnecessary increase in runtime when we increase region count. Co-authored-by: Zac Mrowicki Co-authored-by: Tom Kirchner --- tools/pubsys/src/aws/ami/mod.rs | 36 ++++++++++++++++++++++----------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/tools/pubsys/src/aws/ami/mod.rs b/tools/pubsys/src/aws/ami/mod.rs index 9b6c610d33f..c7ce0e95780 100644 --- a/tools/pubsys/src/aws/ami/mod.rs +++ b/tools/pubsys/src/aws/ami/mod.rs @@ -278,17 +278,27 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result> ec2_clients.insert(region.clone(), ec2_client); } - let mut copy_requests = Vec::with_capacity(regions.len()); + // First, we check if the AMI already exists in each region. + let mut get_requests = Vec::with_capacity(regions.len()); for region in regions.iter() { let ec2_client = &ec2_clients[region]; - if let Some(id) = get_ami_id(&ami_args.name, &ami_args.arch, region.name(), ec2_client) - .await - .context(error::GetAmiId { - name: &ami_args.name, - arch: &ami_args.arch, - region: region.name(), - })? - { + let get_request = get_ami_id(&ami_args.name, &ami_args.arch, region.name(), ec2_client); + let info_future = ready(region.clone()); + get_requests.push(join(info_future, get_request)); + } + let request_stream = stream::iter(get_requests).buffer_unordered(4); + let get_responses: Vec<(Region, std::result::Result, register::Error>)> = + request_stream.collect().await; + + // If an AMI already existed, just add it to our list, otherwise prepare a copy request. + let mut copy_requests = Vec::with_capacity(regions.len()); + for (region, get_response) in get_responses { + let get_response = get_response.context(error::GetAmiId { + name: &ami_args.name, + arch: &ami_args.arch, + region: region.name(), + })?; + if let Some(id) = get_response { info!( "Found '{}' already registered in {}: {}", ami_args.name, @@ -298,14 +308,16 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result> amis.insert(region.name().to_string(), Image::new(&id, &ami_args.name)); continue; } - let request = CopyImageRequest { + + let ec2_client = &ec2_clients[®ion]; + let copy_request = CopyImageRequest { description: ami_args.description.clone(), name: ami_args.name.clone(), source_image_id: ids_of_image.image_id.clone(), source_region: base_region.name().to_string(), ..Default::default() }; - let response_future = ec2_client.copy_image(request); + let copy_future = ec2_client.copy_image(copy_request); let base_region_name = base_region.name(); // Store the region so we can output it to the user @@ -318,7 +330,7 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result> region.name() ) }); - copy_requests.push(message_future.then(|_| join(region_future, response_future))); + copy_requests.push(message_future.then(|_| join(region_future, copy_future))); } // If all target regions already have the AMI, we're done. From aac917ba30d070376bee89a65e65d3c55a4da57a Mon Sep 17 00:00:00 2001 From: Tom Kirchner Date: Thu, 27 Aug 2020 12:02:37 -0700 Subject: [PATCH 2/2] Add info logging before service requests to help explain timing Co-authored-by: Zac Mrowicki Co-authored-by: Tom Kirchner --- tools/pubsys/src/aws/ami/mod.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tools/pubsys/src/aws/ami/mod.rs b/tools/pubsys/src/aws/ami/mod.rs index c7ce0e95780..4e65edf0b4a 100644 --- a/tools/pubsys/src/aws/ami/mod.rs +++ b/tools/pubsys/src/aws/ami/mod.rs @@ -212,6 +212,7 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result> // First we need to find the account IDs for any given roles, so we can grant access to those // accounts to copy the AMI and snapshots. + info!("Getting account IDs for target regions so we can grant access to copy source AMI"); let mut account_ids = get_account_ids(®ions, &base_region, &aws).await?; // Get the account ID used in the base region; we don't need to grant to it so we can remove it @@ -235,6 +236,7 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result> // If we have any accounts other than the base account, grant them access. if !account_ids.is_empty() { + info!("Granting access to target accounts so we can copy the AMI"); let account_id_vec: Vec<_> = account_ids.into_iter().collect(); modify_snapshots( @@ -279,6 +281,7 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result> } // First, we check if the AMI already exists in each region. + info!("Checking whether AMIs already exist in target regions"); let mut get_requests = Vec::with_capacity(regions.len()); for region in regions.iter() { let ec2_client = &ec2_clients[region];