Also prune SSM documents from ec2-test-framework (#1925)

There were some leaking SSM documents that weren't being pruned with the
ec2 instances used. This adds the corresponding logic to remove them.

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license and the ISC license.

tests/ci/cdk/cdk/components.py

@@ -59,6 +59,13 @@ def __init__(self, scope: Construct, id: str, *, project: codebuild.IProject, ec
  "ec2:DescribeInstances",
  ],
  resources=["*"]))
+ lambda_function.add_to_role_policy(
+ iam.PolicyStatement(effect=iam.Effect.ALLOW,
+ actions=[
+ "ssm:ListDocuments",
+ "ssm:DeleteDocument",
+ ],
+ resources=["arn:aws:ssm:{}:{}:*".format(AWS_REGION, AWS_ACCOUNT)]))
 
 
  events.Rule(scope=self, id="PurgeEventRule",

tests/ci/lambda/Cargo.toml

@@ -10,6 +10,7 @@ publish = false
 aws-config = "1.1.7"
 aws-sdk-codebuild = "1.44.0"
 aws-sdk-ec2 = "1.59.0"
+aws-sdk-ssm = "1.51.0"
 aws-sdk-secretsmanager = "1.39.0"
 lambda_runtime = "0.8.0"
 log = "0.4.17"

tests/ci/lambda/src/bin/purge-stale-builds.rs

@@ -5,6 +5,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use aws_sdk_codebuild::types::BuildBatchFilter;
 use aws_sdk_ec2::operation::describe_instances::DescribeInstancesOutput;
 use aws_sdk_ec2::types::Filter;
+use aws_sdk_ssm::types::DocumentKeyValuesFilter;
 use lambda_runtime::{service_fn, Error, LambdaEvent};
 use serde_json::{json, Value};
 
@@ -33,7 +34,7 @@ async fn handle(_event: LambdaEvent<Value>) -> Result<(), Error> {
  std::env::var("GITHUB_TOKEN_SECRET_NAME")
  .map_err(|_| "failed to find github access token secret name")?,
  )
- .await?;
+  .await?;
 
  let github = octocrab::initialise(octocrab::Octocrab::builder().personal_token(github_token))
  .map_err(|e| format!("failed to build github client: {e}"))?;
@@ -42,17 +43,15 @@ async fn handle(_event: LambdaEvent<Value>) -> Result<(), Error> {
 
  let project = std::env::var("CODEBUILD_PROJECT_NAME").unwrap();
 
- let have_ec2_permissions: bool = project == "aws-lc-ci-ec2-test-framework";
+ let is_ec2_test_framework: bool = project == "aws-lc-ci-ec2-test-framework";
 
  let github_repo_owner = std::env::var("GITHUB_REPO_OWNER").unwrap();
 
  let mut pull_requests: HashMap<u64, Vec<CommitBuild>> = HashMap::new();
 
  log::info!("Pulling builds for {project}");
 
- let builds = get_project_build_batches(&codebuild_client, project.clone())
- .await
- .unwrap();
+ let builds = get_project_build_batches(&codebuild_client, project.clone()).await?;
 
  let project_pull_requests = gather_pull_request_builds(&codebuild_client, builds).await?;
 
@@ -67,12 +66,18 @@ async fn handle(_event: LambdaEvent<Value>) -> Result<(), Error> {
  .or_insert(builds);
  }
 
- // Filters for aws-lc-ec2-test-framework specific hosts.
- let ec2_client_optional: Option<aws_sdk_ec2::Client> = if have_ec2_permissions {
+ let ec2_client_optional: Option<aws_sdk_ec2::Client> = if is_ec2_test_framework {
  Some(aws_sdk_ec2::Client::new(&config))
  } else {
  None
  };
+ let ssm_client_optional: Option<aws_sdk_ssm::Client> = if is_ec2_test_framework {
+ Some(aws_sdk_ssm::Client::new(&config))
+ } else {
+ None
+ };
+
+ // Filters for aws-lc-ec2-test-framework specific hosts.
  let ec2_framework_filters = vec![
  Filter::builder()
  .name("instance-state-name")
@@ -106,6 +111,7 @@ async fn handle(_event: LambdaEvent<Value>) -> Result<(), Error> {
  None
  };
 
+ let mut ssm_deleted_documents: Vec<String> = vec![];
  let mut ec2_terminated_instances: Vec<String> = vec![];
  let mut stopped_builds: u64 = 0;
 
@@ -137,6 +143,7 @@ async fn handle(_event: LambdaEvent<Value>) -> Result<(), Error> {
  .map_err(|e| format!("failed to stop_build_batch: {e}"))?;
  stopped_builds += 1;
 
+ // Gather a list of all unneeded ec2 instances and terminate after the loop.
  if let Some(ref ec2_describe_response) = ec2_describe_response_optional {
  // Prune unneeded ec2 instances.
  for reservation in ec2_describe_response.reservations() {
@@ -154,6 +161,10 @@ async fn handle(_event: LambdaEvent<Value>) -> Result<(), Error> {
  }
  }
  }
+
+ // Gather a list of old commits. The SSM documents should have the commits within
+ // their document name.
+ ssm_deleted_documents.push(old_commit.to_string());
  }
  }
 
@@ -169,6 +180,28 @@ async fn handle(_event: LambdaEvent<Value>) -> Result<(), Error> {
  }
  }
 
+ if !ssm_deleted_documents.is_empty() && is_ec2_test_framework {
+ log::info!("Query for list of documents to delete with: {:?}",ssm_deleted_documents);
+
+ let all_documents = get_ssm_document_list(
+ &ssm_client_optional,
+ ssm_deleted_documents.clone(),
+ ).await?;
+
+ // Prune hanging ssm documents corresponding to commits.
+ for document in all_documents {
+ log::info!("SSM document {:?} will be deleted", document);
+ if let Some(ref ssm_client) = ssm_client_optional {
+ ssm_client
+ .delete_document()
+ .name(document)
+ .send()
+ .await
+ .map_err(|e| format!("failed to delete ssm document: {e}"))?;
+ }
+ }
+ }
+
  let timestamp = SystemTime::now()
  .duration_since(UNIX_EPOCH)
  .unwrap()
@@ -337,3 +370,46 @@ async fn get_project_build_batches(
 
  Ok(builds)
 }
+
+async fn get_ssm_document_list(
+ client: &Option<aws_sdk_ssm::Client>,
+ ssm_deleted_documents: Vec<String>,
+) -> Result<Vec<String>, String> {
+ let ssm_client_filters = vec![
+ DocumentKeyValuesFilter::builder()
+ .key("SearchKeyword")
+ .set_values(Some(ssm_deleted_documents))
+ .build(),
+ DocumentKeyValuesFilter::builder()
+ .key("Owner")
+ .values("Self")
+ .build(),
+ ];
+ log::info!("Document filter list: {:?}", ssm_client_filters);
+
+ let mut document_list: Vec<String> = vec![];
+
+ if let Some(ref ssm_client) = client {
+ let mut paginator = ssm_client
+ .list_documents()
+ .set_filters(Some(ssm_client_filters))
+ .into_paginator()
+ .send();
+
+ while let Some(result) = paginator.next().await {
+ if result.is_err() {
+ return Err(format!("SSM ListDocuments Failed: {}", result.unwrap_err()));
+ }
+
+ let document_ids = Vec::from(result.unwrap().document_identifiers());
+
+ for document in document_ids {
+ document_list.push(document.name.unwrap())
+ }
+ }
+ }
+
+ log::info!("Found documents to delete {:?}", document_list);
+
+ Ok(document_list)
+}