fix: ensure abe2e SIG image replication (#4589)

Co-authored-by: Cameron Meissner <[email protected]>

e2e/config/vhd.go

@@ -4,96 +4,158 @@ import (
 	"context"
 	"fmt"
 	"log"
-	"regexp"
 	"strings"
 	"sync"
 	"time"
 
-	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute"
 )
 
 const (
 	imageGallery       = "/subscriptions/8ecadfc9-d1a3-4ea4-b844-0d9f87e4d7c8/resourceGroups/aksvhdtestbuildrg/providers/Microsoft.Compute/galleries/PackerSigGalleryEastUS/images/"
 	noSelectionTagName = "abe2e-ignore"
+
+	fetchResourceIDTimeout = 3 * time.Minute
 )
 
 var (
-	VHDUbuntu1804Gen2Containerd               = newVHDResourceIDFetcher(imageGallery + "1804Gen2")
-	VHDUbuntu2204Gen2Arm64Containerd          = newVHDResourceIDFetcher(imageGallery + "2204Gen2Arm64")
-	VHDUbuntu2204Gen2Containerd               = newVHDResourceIDFetcher(imageGallery + "2204Gen2")
-	VHDAzureLinuxV2Gen2Arm64                  = newVHDResourceIDFetcher(imageGallery + "AzureLinuxV2Gen2Arm64")
-	VHDAzureLinuxV2Gen2                       = newVHDResourceIDFetcher(imageGallery + "AzureLinuxV2Gen2")
-	VHDCBLMarinerV2Gen2Arm64                  = newVHDResourceIDFetcher(imageGallery + "CBLMarinerV2Gen2Arm64")
-	VHDCBLMarinerV2Gen2                       = newVHDResourceIDFetcher(imageGallery + "CBLMarinerV2Gen2")
-	VHDUbuntu2204Gen2ContainerdPrivateKubePkg = func() (VHDResourceID, error) {
-		// this is a particular 2204Gen2 image originally built with private packages,
-		// if we ever want to update this then we'd need to run a new VHD build using private package overrides
-		return imageGallery + "2204Gen2/versions/1.1704411049.2812", nil
-	}
+	VHDUbuntu1804Gen2Containerd      = newSIGImageVersionResourceIDFetcher(imageGallery + "1804Gen2")
+	VHDUbuntu2204Gen2Arm64Containerd = newSIGImageVersionResourceIDFetcher(imageGallery + "2204Gen2Arm64")
+	VHDUbuntu2204Gen2Containerd      = newSIGImageVersionResourceIDFetcher(imageGallery + "2204Gen2")
+	VHDAzureLinuxV2Gen2Arm64         = newSIGImageVersionResourceIDFetcher(imageGallery + "AzureLinuxV2Gen2Arm64")
+	VHDAzureLinuxV2Gen2              = newSIGImageVersionResourceIDFetcher(imageGallery + "AzureLinuxV2Gen2")
+	VHDCBLMarinerV2Gen2Arm64         = newSIGImageVersionResourceIDFetcher(imageGallery + "CBLMarinerV2Gen2Arm64")
+	VHDCBLMarinerV2Gen2              = newSIGImageVersionResourceIDFetcher(imageGallery + "CBLMarinerV2Gen2")
+
+	// this is a particular 2204Gen2 image originally built with private packages,
+	// if we ever want to update this then we'd need to run a new VHD build using private package overrides
+	VHDUbuntu2204Gen2ContainerdPrivateKubePkg = newStaticSIGImageVersionResourceIDFetcher(imageGallery + "2204Gen2/versions/1.1704411049.2812")
 )
 
+var ErrNotFound = fmt.Errorf("not found")
+
+type sigImageDefinition struct {
+	subscriptionID string
+	resourceGroup  string
+	gallery        string
+	definition     string
+}
+
+type sigImageVersion struct {
+	sigImageDefinition
+	version string
+}
+
+func newSIGImageDefinitionFromResourceID(resourceID *arm.ResourceID) sigImageDefinition {
+	return sigImageDefinition{
+		subscriptionID: resourceID.SubscriptionID,
+		resourceGroup:  resourceID.ResourceGroupName,
+		gallery:        resourceID.Parent.Name,
+		definition:     resourceID.Name,
+	}
+}
+
+func newSIGImageVersionFromResourceID(resourceID *arm.ResourceID) sigImageVersion {
+	return sigImageVersion{
+		sigImageDefinition: newSIGImageDefinitionFromResourceID(resourceID.Parent),
+		version:            resourceID.Name,
+	}
+}
+
 // VHDResourceID represents a resource ID pointing to a VHD in Azure. This could be theoretically
 // be the resource ID of a managed image or SIG image version, though for now this will always be a SIG image version.
 type VHDResourceID string
 
-// newVHDResourceIDFetcher is a factory function
+func (id VHDResourceID) Short() string {
+	sep := "Microsoft.Compute/galleries/"
+	str := string(id)
+	if strings.Contains(str, sep) && !strings.HasSuffix(str, sep) {
+		return strings.Split(str, sep)[1]
+	}
+	return str
+}
+
+// newSIGImageVersionResourceIDFetcher is a factory function
 // it returns a function that fetches the latest VHDResourceID for a given image
 // the function is memoized and will only evaluate once on the first call
-func newVHDResourceIDFetcher(image string) func() (VHDResourceID, error) {
+func newSIGImageVersionResourceIDFetcher(imageDefinitionResourceID string) func() (VHDResourceID, error) {
 	resourceID := VHDResourceID("")
 	var err error
 	once := sync.Once{}
 	// evaluate the function once and cache the result
 	return func() (VHDResourceID, error) {
 		once.Do(func() {
-			resourceID, err = findLatestImageWithTag(image, SIGVersionTagName, SIGVersionTagValue)
+			resourceID, err = findLatestSIGImageVersionWithTag(imageDefinitionResourceID, SIGVersionTagName, SIGVersionTagValue)
 			if err != nil {
-				err = fmt.Errorf("img: %s, tag %s=%s", image, SIGVersionTagName, SIGVersionTagValue)
+				err = fmt.Errorf("img: %s, tag %s=%s, err %w", imageDefinitionResourceID, SIGVersionTagName, SIGVersionTagValue, err)
 				log.Printf("failed to find the latest image %s", err)
 			} else {
-				log.Printf("Resource ID for %s: %s", image, resourceID)
+				log.Printf("Resource ID for %s: %s", imageDefinitionResourceID, resourceID)
 			}
 		})
 		return resourceID, err
 	}
 }
 
-func (id VHDResourceID) Short() string {
-	sep := "Microsoft.Compute/galleries/"
-	str := string(id)
-	if strings.Contains(str, sep) && !strings.HasSuffix(str, sep) {
-		return strings.Split(str, sep)[1]
+func newStaticSIGImageVersionResourceIDFetcher(imageVersionResourceID string) func() (VHDResourceID, error) {
+	resourceID := VHDResourceID("")
+	var err error
+	once := sync.Once{}
+
+	return func() (VHDResourceID, error) {
+		once.Do(func() {
+			resourceID, err = ensureStaticSIGImageVersion(imageVersionResourceID)
+			if err != nil {
+				err = fmt.Errorf("img: %s, err: %w", imageVersionResourceID, err)
+				log.Printf("failed to find static image %s", err)
+			} else {
+				log.Printf("Resource ID for %s: %s", imageVersionResourceID, resourceID)
+			}
+		})
+		return resourceID, err
 	}
-	return str
 }
 
-var ErrNotFound = fmt.Errorf("not found")
-
-func findLatestImageWithTag(imageID, tagName, tagValue string) (VHDResourceID, error) {
-	ctx, cancel := context.WithTimeout(context.TODO(), time.Minute)
+func ensureStaticSIGImageVersion(imageVersionResourceID string) (VHDResourceID, error) {
+	ctx, cancel := context.WithTimeout(context.TODO(), fetchResourceIDTimeout)
 	defer cancel()
-	image, err := parseImageID(imageID)
+
+	rid, err := arm.ParseResourceID(imageVersionResourceID)
 	if err != nil {
-		return "", err
+		return "", fmt.Errorf("parsing image version resouce ID: %w", err)
 	}
+	version := newSIGImageVersionFromResourceID(rid)
 
-	cred, err := azidentity.NewDefaultAzureCredential(nil)
+	resp, err := Azure.GalleryImageVersionClient.Get(ctx, version.resourceGroup, version.gallery, version.definition, version.version, nil)
 	if err != nil {
-		return "", fmt.Errorf("failed to obtain a credential: %v", err)
+		return "", fmt.Errorf("getting live image version info: %w", err)
+	}
+
+	if err := ensureReplication(ctx, version.sigImageDefinition, &resp.GalleryImageVersion); err != nil {
+		return "", fmt.Errorf("ensuring image replication: %w", err)
 	}
 
-	client, err := armcompute.NewGalleryImageVersionsClient(image.subscriptionID, cred, nil)
+	return VHDResourceID(imageVersionResourceID), nil
+}
+
+func findLatestSIGImageVersionWithTag(imageDefinitionResourceID, tagName, tagValue string) (VHDResourceID, error) {
+	ctx, cancel := context.WithTimeout(context.TODO(), fetchResourceIDTimeout)
+	defer cancel()
+
+	rid, err := arm.ParseResourceID(imageDefinitionResourceID)
 	if err != nil {
-		return "", fmt.Errorf("failed to create a new images client: %v", err)
+		return "", fmt.Errorf("parsing image definition resource ID: %w", err)
 	}
+	definition := newSIGImageDefinitionFromResourceID(rid)
 
-	pager := client.NewListByGalleryImagePager(image.resourceGroup, image.galleryName, image.imageName, nil)
+	pager := Azure.GalleryImageVersionClient.NewListByGalleryImagePager(definition.resourceGroup, definition.gallery, definition.definition, nil)
 	var latestVersion *armcompute.GalleryImageVersion
 	for pager.More() {
 		page, err := pager.NextPage(ctx)
 		if err != nil {
-			return "", fmt.Errorf("failed to get next page: %v", err)
+			return "", fmt.Errorf("failed to get next page: %w", err)
 		}
 		versions := page.Value
 		for _, version := range versions {
@@ -114,29 +176,46 @@ func findLatestImageWithTag(imageID, tagName, tagValue string) (VHDResourceID, e
 	if latestVersion == nil {
 		return "", ErrNotFound
 	}
+
+	if err := ensureReplication(ctx, definition, latestVersion); err != nil {
+		return "", fmt.Errorf("ensuring image replication: %w", err)
+	}
+
 	return VHDResourceID(*latestVersion.ID), nil
 }
 
-type imageID struct {
-	subscriptionID string
-	resourceGroup  string
-	galleryName    string
-	imageName      string
+func ensureReplication(ctx context.Context, definition sigImageDefinition, version *armcompute.GalleryImageVersion) error {
+	if replicatedToCurrentRegion(version) {
+		return nil
+	}
+	return replicateToCurrentRegion(ctx, definition, version)
+}
+
+func replicatedToCurrentRegion(version *armcompute.GalleryImageVersion) bool {
+	for _, targetRegion := range version.Properties.PublishingProfile.TargetRegions {
+		if strings.EqualFold(strings.ReplaceAll(*targetRegion.Name, " ", ""), Location) {
+			return true
+		}
+	}
+	return false
 }
 
-func parseImageID(resourceID string) (imageID, error) {
-	pattern := `(?i)^/subscriptions/([^/]+)/resourceGroups/([^/]+)/providers/Microsoft\.Compute/galleries/([^/]+)/images/([^/]+)$`
-	re := regexp.MustCompile(pattern)
-	matches := re.FindStringSubmatch(resourceID)
+func replicateToCurrentRegion(ctx context.Context, definition sigImageDefinition, version *armcompute.GalleryImageVersion) error {
+	log.Printf("will replicate image version %s to region %s...", *version.ID, Location)
+
+	version.Properties.PublishingProfile.TargetRegions = append(version.Properties.PublishingProfile.TargetRegions, &armcompute.TargetRegion{
+		Name:                 &Location,
+		RegionalReplicaCount: to.Ptr[int32](1),
+		StorageAccountType:   to.Ptr(armcompute.StorageAccountTypeStandardLRS),
+	})
 
-	if matches == nil || len(matches) != 5 {
-		return imageID{}, fmt.Errorf("failed to parse image ID %q", resourceID)
+	resp, err := Azure.GalleryImageVersionClient.BeginCreateOrUpdate(ctx, definition.resourceGroup, definition.gallery, definition.definition, *version.Name, *version, nil)
+	if err != nil {
+		return fmt.Errorf("begin updating image version target regions: %w", err)
+	}
+	if _, err := resp.PollUntilDone(ctx, nil); err != nil {
+		return fmt.Errorf("updating image version target regions: %w", err)
 	}
 
-	return imageID{
-		subscriptionID: matches[1],
-		resourceGroup:  matches[2],
-		galleryName:    matches[3],
-		imageName:      matches[4],
-	}, nil
+	return nil
 }