Also prune SSM documents from ec2-test-framework (aws#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)
+}