From 5468150cc842213a5f4f4a1952ce91922ed0ab12 Mon Sep 17 00:00:00 2001
From: Cameron Meissner <cameronmeissner@gmail.com>
Date: Tue, 3 Dec 2024 12:01:24 -0800
Subject: [PATCH] feat: dedicated garbage collection pipeline (#5325)

Co-authored-by: Cameron Meissner <cameissner@microsoft.com>
---
 .pipelines/.vsts-garabge-collection.yaml | 24 ++++++++++
 vhdbuilder/packer/test/run-test.sh       |  2 +-
 vhdbuilder/packer/vhd-scanning.sh        |  2 +-
 vhdbuilder/scripts/gc.sh                 | 61 ++++++++++++++++++++++++
 4 files changed, 87 insertions(+), 2 deletions(-)
 create mode 100644 .pipelines/.vsts-garabge-collection.yaml
 create mode 100644 vhdbuilder/scripts/gc.sh

diff --git a/.pipelines/.vsts-garabge-collection.yaml b/.pipelines/.vsts-garabge-collection.yaml
new file mode 100644
index 00000000000..2d04f3fae35
--- /dev/null
+++ b/.pipelines/.vsts-garabge-collection.yaml
@@ -0,0 +1,24 @@
+pool:
+  name: $(POOL_NAME)
+
+parameters:
+- name: DRY_RUN
+  displayName: Dry Run
+  type: boolean
+  default: false
+
+jobs:
+- job: gc
+  displayName: Garbage Collection
+  steps:
+    - checkout: self
+      fetchTags: false
+      fetchDepth: 1
+      
+    - bash: |
+        chmod +x ./vhdbuilder/scripts/gc.sh
+        ./vhdbuilder/scripts/gc.sh
+      env:
+        SUBSCRIPTION_ID: $(SUBSCRIPTION_ID)
+        DRY_RUN: ${{ parameters.DRY_RUN }}
+      displayName: Garbage collect resource groups
diff --git a/vhdbuilder/packer/test/run-test.sh b/vhdbuilder/packer/test/run-test.sh
index 94959456425..0a29cdff3cd 100755
--- a/vhdbuilder/packer/test/run-test.sh
+++ b/vhdbuilder/packer/test/run-test.sh
@@ -32,7 +32,7 @@ else
     exit 1
   fi
 fi
-az group create --name $TEST_VM_RESOURCE_GROUP_NAME --location ${AZURE_LOCATION} --tags "source=AgentBaker" "branch=${GIT_BRANCH}"
+az group create --name $TEST_VM_RESOURCE_GROUP_NAME --location ${AZURE_LOCATION} --tags "source=AgentBaker,now=$(date +%s)" "branch=${GIT_BRANCH}"
 
 # defer function to cleanup resource group when VHD debug is not enabled
 function cleanup() {
diff --git a/vhdbuilder/packer/vhd-scanning.sh b/vhdbuilder/packer/vhd-scanning.sh
index aa7d0934688..47d4499f6cd 100755
--- a/vhdbuilder/packer/vhd-scanning.sh
+++ b/vhdbuilder/packer/vhd-scanning.sh
@@ -42,7 +42,7 @@ SCAN_VM_ADMIN_PASSWORD="ScanVM@$(date +%s)"
 set -x
 
 RESOURCE_GROUP_NAME="$SCAN_RESOURCE_PREFIX-$(date +%s)-$RANDOM"
-az group create --name $RESOURCE_GROUP_NAME --location ${PACKER_BUILD_LOCATION} --tags "source=AgentBaker" "branch=${GIT_BRANCH}"
+az group create --name $RESOURCE_GROUP_NAME --location ${PACKER_BUILD_LOCATION} --tags "source=AgentBaker,now=$(date +%s)" "branch=${GIT_BRANCH}"
 
 function cleanup() {
     echo "Deleting resource group ${RESOURCE_GROUP_NAME}"
diff --git a/vhdbuilder/scripts/gc.sh b/vhdbuilder/scripts/gc.sh
new file mode 100644
index 00000000000..0c367afe2aa
--- /dev/null
+++ b/vhdbuilder/scripts/gc.sh
@@ -0,0 +1,61 @@
+#!/bin/bash
+set -euxo pipefail
+
+[ -z "${SUBSCRIPTION_ID:-}" ] && echo "SUBSCRIPTION_ID must be set" && exit 1
+
+SKIP_TAG_NAME="gc.skip"
+SKIP_TAG_VALUE="true"
+
+DRY_RUN="${DRY_RUN:-}"
+
+DAY_AGO=$(( $(date +%s) - 86400 )) # 24 hours ago
+WEEK_AGO=$(( $(date +%s) - 604800 )) # 7 days ago
+
+function main() {
+    az login --identity # relies on an appropriately permissioned identity being attached to the build agent
+    az account set -s $SUBSCRIPTION_ID
+
+    echo "garbage collecting ephemeral resource groups..."
+    cleanup_rgs || exit $?
+
+    # TODO(cameissner): migrate linux VHD build back-fill deletion logic to this script
+}
+
+function cleanup_rgs() {
+    groups=$(az group list | jq -r --arg dl $DAY_AGO '.[] | select(.name | test("vhd-test*|vhd-scanning*|pkr-Resource-Group*")) | select(.tags.now < $dl).name'  | tr -d '\"' || "")
+    if [ -z "$groups" ]; then
+        echo "no resource groups found for garbage collection"
+        return 0
+    fi
+
+    for group in $groups; do
+        echo "resource group $group is in-scope for garbage collection"
+        group_object=$(az group show -g $group)
+        tag_value=$(echo "$group_object" | jq -r --arg skipTagName $SKIP_TAG_NAME '.tags."\($skipTagName)"')
+
+        if [ "${tag_value,,}" == "$SKIP_TAG_VALUE" ]; then
+            now=$(echo "$group_object" | jq -r '.tags.now')
+            if [ "$now" != "null" ] && [ $now -lt $WEEK_AGO ]; then
+                echo "resource group $group is tagged with $SKIP_TAG_NAME=$SKIP_TAG_VALUE but is more than 7 days old, will attempt to delete..."
+                delete_group $group || return $?
+            fi
+            continue
+        fi
+
+        echo "will attempt to delete resource group $group"
+        delete_group $group || return $?
+    done
+}
+
+function delete_group() {
+    local group=$1
+
+    if [ "${DRY_RUN,,}" == "true" ]; then
+        echo "DRY_RUN: az group delete -g $group --yes --no-wait"
+        return 0
+    fi
+
+    az group delete -g $group --yes --no-wait || return $?
+}
+
+main "$@"
\ No newline at end of file