From 3b0b34e3e279a623d4c9495205d409141dcb691c Mon Sep 17 00:00:00 2001 From: Varun <48163435+varunch77@users.noreply.github.com> Date: Fri, 4 Oct 2024 11:47:45 -0400 Subject: [PATCH] Update AMI cleaner script to properly handle macOS images (#1370) --- tool/clean/clean_ami/clean_ami.go | 134 +++++++++++++++++++++++++----- 1 file changed, 115 insertions(+), 19 deletions(-) diff --git a/tool/clean/clean_ami/clean_ami.go b/tool/clean/clean_ami/clean_ami.go index 82ea91068a..edfdc9c7b0 100644 --- a/tool/clean/clean_ami/clean_ami.go +++ b/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) }