Update AMI cleaner script to properly handle macOS images (aws#1370)

Author: varunch77

tool/clean/clean_ami/clean_ami.go

@@ -8,8 +8,11 @@ package main
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
+	"sort"
+	"strings"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -22,18 +25,113 @@ import (
 )
 
 func main() {
-	err := cleanAMI()
+	err := cleanAMIs()
 	if err != nil {
 		log.Fatalf("errors cleaning %v", err)
 	}
 }
 
-func cleanAMI() error {
+// takes a list of AMIs and sorts them by creation date (youngest to oldest)
+func sortAMIsByCreationDate(amiList []types.Image, errList *[]error) []types.Image {
+	sort.Slice(amiList, func(i, j int) bool {
+		if amiList[i].CreationDate != nil && amiList[j].CreationDate != nil {
+			iCreationDate, iErr := smithyTime.ParseDateTime(*amiList[i].CreationDate)
+			jCreationDate, jErr := smithyTime.ParseDateTime(*amiList[j].CreationDate)
+
+			if err := errors.Join(iErr, jErr); err != nil && errList != nil {
+				*errList = append(*errList, err)
+				return false
+			}
+
+			return iCreationDate.After(jCreationDate)
+		} else {
+			return false
+		}
+	})
+
+	return amiList
+}
+
+// given a slice of AMIs, deregisters them one by one
+func deregisterAMIs(ctx context.Context, ec2client *ec2.Client, images []types.Image, errList *[]error) {
+	for _, image := range images {
+		if image.Name != nil && image.ImageId != nil && image.CreationDate != nil {
+			log.Printf("Try to delete ami %v tags %v image id %v image creation date raw %v", *image.Name, image.Tags, *image.ImageId, *image.CreationDate)
+			deregisterImageInput := &ec2.DeregisterImageInput{ImageId: image.ImageId}
+			_, err := ec2client.DeregisterImage(ctx, deregisterImageInput)
+
+			if err != nil && errList != nil {
+				log.Printf("Error while deregistering ami %v", *image.Name)
+				*errList = append(*errList, err)
+			}
+		}
+	}
+}
+
+// given a map of macos version/architecture to a list of corresponding AMIs, deregister AMIs that are no longer needed
+func cleanMacAMIs(ctx context.Context, ec2client *ec2.Client, macosImageAmiMap map[string][]types.Image, expirationDate time.Time, errList *[]error) {
+	for name, amiList := range macosImageAmiMap {
+		// don't delete an ami if it's the only one for that version/architecture
+		if len(amiList) == 1 {
+			continue
+		}
+
+		// Sort AMIs by creation date (youngest to oldest)
+		amiList = sortAMIsByCreationDate(amiList, errList)
+
+		// find the youngest AMI in the list
+		youngestCreationDate, err := smithyTime.ParseDateTime(aws.ToString(amiList[0].CreationDate))
+
+		if err != nil && errList != nil {
+			*errList = append(*errList, err)
+			continue
+		}
+
+		if expirationDate.After(youngestCreationDate) {
+			// If the youngest AMI is over 60 days old, we keep one (the youngest) and can delete the rest
+			log.Printf("Youngest AMI for %s is over 60 days old. Deleting all but the youngest.", name)
+			deregisterAMIs(ctx, ec2client, amiList[1:], errList)
+		} else {
+			// If the youngest AMI is under 60 days old, keep incrementing until we find AMIs older than 60 days and delete them
+			for index, ami := range amiList {
+				creationDate, err := smithyTime.ParseDateTime(aws.ToString(ami.CreationDate))
+				if err != nil && errList != nil {
+					*errList = append(*errList, err)
+					continue
+				}
+				if expirationDate.After(creationDate) {
+					// once you find the first AMI that's over 60 days old, delete the ones that follow
+					deregisterAMIs(ctx, ec2client, amiList[index:], errList)
+					break
+				}
+			}
+		}
+	}
+}
+
+// given a single non macos image, determine its age and deregister if needed
+func cleanNonMacAMIs(ctx context.Context, ec2client *ec2.Client, image types.Image, expirationDate time.Time, errList *[]error) {
+	creationDate, err := smithyTime.ParseDateTime(aws.ToString(image.CreationDate))
+	if err != nil && errList != nil {
+		*errList = append(*errList, err)
+		return
+	}
+
+	if expirationDate.After(creationDate) {
+		deregisterAMIs(ctx, ec2client, []types.Image{image}, errList)
+	}
+}
+
+func cleanAMIs() error {
 	log.Print("Begin to clean EC2 AMI")
 
+	// sets expiration date to 60 days in the past
 	expirationDate := time.Now().UTC().Add(clean.KeepDurationSixtyDay)
-	cxt := context.Background()
-	defaultConfig, err := config.LoadDefaultConfig(cxt)
+	log.Printf("Expiration date set as %v", expirationDate)
+
+	// load default config
+	ctx := context.Background()
+	defaultConfig, err := config.LoadDefaultConfig(ctx)
 	if err != nil {
 		return err
 	}
@@ -46,30 +144,28 @@ func cleanAMI() error {
 
 	//get instances to delete
 	describeImagesInput := ec2.DescribeImagesInput{Filters: []types.Filter{nameFilter}}
-	describeImagesOutput, err := ec2client.DescribeImages(cxt, &describeImagesInput)
+	describeImagesOutput, err := ec2client.DescribeImages(ctx, &describeImagesInput)
 	if err != nil {
 		return err
 	}
 
 	var errList []error
+	// stores a list of AMIs per each macos version/architecture
+	macosImageAmiMap := make(map[string][]types.Image)
+
 	for _, image := range describeImagesOutput.Images {
-		creationDate, err := smithyTime.ParseDateTime(*image.CreationDate)
-		if err != nil {
-			errList = append(errList, err)
-			continue
-		}
-		log.Printf("image name %v image id %v experation date %v creation date parsed %v image creation date raw %v",
-			*image.Name, *image.ImageId, creationDate, expirationDate, *image.CreationDate)
-		if expirationDate.After(creationDate) {
-			log.Printf("Try to delete ami %s tags %v launch-date %s", *image.Name, image.Tags, *image.CreationDate)
-			deregisterImageInput := ec2.DeregisterImageInput{ImageId: image.ImageId}
-			_, err := ec2client.DeregisterImage(cxt, &deregisterImageInput)
-			if err != nil {
-				errList = append(errList, err)
-			}
+		if image.Name != nil && strings.HasPrefix(*image.Name, "cloudwatch-agent-integration-test-mac") {
+			// mac image - add it to the map and do nothing else for now
+			macosImageAmiMap[*image.Name] = append(macosImageAmiMap[*image.Name], image)
+		} else {
+			// non mac image - clean it if it's older than 60 days
+			cleanNonMacAMIs(ctx, ec2client, image, expirationDate, &errList)
 		}
 	}
 
+	// handle the mac AMIs
+	cleanMacAMIs(ctx, ec2client, macosImageAmiMap, expirationDate, &errList)
+
 	if len(errList) != 0 {
 		return fmt.Errorf("%v", errList)
 	}