fix(e2e): replicate gallery images to a fixed region set

.github/workflows/validate-e2e-regions.yml

@@ -0,0 +1,20 @@
+name: validate-e2e-regions
+on: pull_request
+
+# The e2e module is not covered by `make test`, which runs `go test ./...` from the repo root
+# and does not cross module boundaries. Without this the region checks would never run in CI.
+jobs:
+  validate-e2e-regions:
+    runs-on: ubuntu-24.04
+    permissions:
+      contents: read
+    steps:
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      with:
+        persist-credentials: false
+    - uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7
+      with:
+        go-version-file: e2e/go.mod
+    - run: |
+        cd e2e && go test ./config/...
+      name: Verify the gallery replication region list matches the scenarios

e2e/config/azure.go

@@ -652,30 +652,99 @@ func (a *AzureClient) LatestSIGImageVersionByTag(ctx context.Context, image *Ima
 	return VHDResourceID(*latestVersion.ID), nil
 }
 
+// ensureReplication makes an image version usable in every region E2E runs in.
+//
+// It replicates to Image.replicationRegions() rather than to the caller's region, which is
+// what stops concurrent writers from clobbering each other. See e2e/config/regions.go.
 func (a *AzureClient) ensureReplication(ctx context.Context, image *Image, version *armcompute.GalleryImageVersion, location string) error {
-	// Wait for any ongoing update operations to complete first
-	if err := a.waitForVersionOperationCompletion(ctx, image, version); err != nil {
-		return fmt.Errorf("waiting for version operation completion: %w", err)
+	desired := image.replicationRegions()
+	if image.Ephemeral {
+		desired = []string{NormalizeRegion(location)}
+	} else if !slices.Contains(desired, NormalizeRegion(location)) {
+		// Appending the caller's region here is what caused the original lost update, so
+		// widening the set is a code change rather than a runtime decision.
+		return fmt.Errorf("image %s is not replicated to %s: add the region to e2eRegions in e2e/config/regions.go, or run E2E in one of %s",
+			image.Name, location, strings.Join(desired, ", "))
+	}
+
+	const maxAttempts = 4
+	for attempt := 1; ; attempt++ {
+		if err := a.waitForVersionOperationCompletion(ctx, image, version); err != nil {
+			return fmt.Errorf("waiting for version operation completion: %w", err)
+		}
+
+		missing := missingRegions(version.Properties.PublishingProfile.TargetRegions, desired)
+		if len(missing) == 0 {
+			toolkit.Logf(ctx, "Image version %s is already replicated to %s", *version.ID, strings.Join(desired, ", "))
+			return nil
+		}
+
+		toolkit.Logf(ctx, "Replicating image version %s to missing regions: %s", *version.ID, strings.Join(missing, ", "))
+		toolkit.Logf(ctx, "##vso[task.logissue type=warning;]Replicating to regions %s", strings.Join(missing, ", "))
+
+		start := time.Now()
+		err := a.replicateImageVersionToRegions(ctx, image, version, missing)
+		toolkit.LogDuration(ctx, time.Since(start), 3*time.Minute, fmt.Sprintf("Replication took: %s (%s)", time.Since(start), *version.ID))
+
+		if err == nil {
+			return nil
+		}
+		if !isGalleryUpdateConflict(err) || attempt >= maxAttempts {
+			return err
+		}
+
+		// Another writer got there first with the same desired state, so there is nothing to
+		// merge. Back off before re-reading: its update is still in flight, and a GET issued
+		// immediately would return the pre-update state and burn the remaining attempts.
+		toolkit.Logf(ctx, "Concurrent update of image version %s; backing off before retry %d/%d", *version.ID, attempt+1, maxAttempts)
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("waiting to retry replication: %w (original error: %v)", ctx.Err(), err)
+		case <-time.After(time.Duration(attempt) * 30 * time.Second):
+		}
+
+		live, getErr := a.getImageVersion(ctx, image, *version.Name)
+		if getErr != nil {
+			return fmt.Errorf("re-reading image version after conflict: %w (original error: %v)", getErr, err)
+		}
+		*version = *live
 	}
+}
 
-	if replicatedToCurrentRegion(version, location) {
-		toolkit.Logf(ctx, "Image version %s is already in region %s", *version.ID, location)
-		return nil
+func (a *AzureClient) getImageVersion(ctx context.Context, image *Image, version string) (*armcompute.GalleryImageVersion, error) {
+	client, err := armcompute.NewGalleryImageVersionsClient(image.Gallery.SubscriptionID, a.Credential, a.ArmOptions)
+	if err != nil {
+		return nil, fmt.Errorf("create image version client: %w", err)
 	}
-	regions := make([]string, 0, len(version.Properties.PublishingProfile.TargetRegions))
-	for _, targetRegion := range version.Properties.PublishingProfile.TargetRegions {
-		regions = append(regions, *targetRegion.Name)
+	resp, err := client.Get(ctx, image.Gallery.ResourceGroupName, image.Gallery.Name, image.Name, version, nil)
+	if err != nil {
+		return nil, fmt.Errorf("get image version %s/%s: %w", image.Name, version, err)
 	}
-	toolkit.Logf(ctx, "Replicating to region %s, available regions: %s, image version %s", location, strings.Join(regions, ", "), *version.ID)
-	toolkit.Logf(ctx, "##vso[task.logissue type=warning;]Replicating to region %s", location)
+	return &resp.GalleryImageVersion, nil
+}
 
-	start := time.Now() // Record the start time
-	err := a.replicateImageVersionToCurrentRegion(ctx, image, version, location)
-	elapsed := time.Since(start) // Calculate the elapsed time
+func isGalleryUpdateConflict(err error) bool {
+	var respErr *azcore.ResponseError
+	return errors.As(err, &respErr) &&
+		(respErr.StatusCode == http.StatusConflict || respErr.StatusCode == http.StatusPreconditionFailed)
+}
 
-	toolkit.LogDuration(ctx, elapsed, 3*time.Minute, fmt.Sprintf("Replication took: %s (%s)", elapsed, *version.ID))
+// missingRegions returns the desired regions that are not already replication targets.
+func missingRegions(existing []*armcompute.TargetRegion, desired []string) []string {
+	present := make(map[string]struct{}, len(existing))
+	for _, region := range existing {
+		if region != nil && region.Name != nil {
+			present[NormalizeRegion(*region.Name)] = struct{}{}
+		}
+	}
 
-	return err
+	var missing []string
+	for _, region := range desired {
+		if _, ok := present[NormalizeRegion(region)]; !ok {
+			missing = append(missing, region)
+		}
+	}
+	return missing
 }
 
 func (a *AzureClient) waitForVersionOperationCompletion(ctx context.Context, image *Image, version *armcompute.GalleryImageVersion) error {
@@ -732,16 +801,21 @@ func (a *AzureClient) waitForVersionOperationCompletion(ctx context.Context, ima
 	return nil
 }
 
-func (a *AzureClient) replicateImageVersionToCurrentRegion(ctx context.Context, image *Image, version *armcompute.GalleryImageVersion, location string) error {
+// replicateImageVersionToRegions appends the missing regions in a single update. Only regions
+// in the fixed set are guaranteed to survive; one added out-of-band by another tool can still
+// be dropped by a stale snapshot, which would need an If-Match merge to prevent.
+func (a *AzureClient) replicateImageVersionToRegions(ctx context.Context, image *Image, version *armcompute.GalleryImageVersion, missing []string) error {
 	galleryImageVersion, err := armcompute.NewGalleryImageVersionsClient(image.Gallery.SubscriptionID, a.Credential, a.ArmOptions)
 	if err != nil {
 		return fmt.Errorf("create a new images client: %v", err)
 	}
-	version.Properties.PublishingProfile.TargetRegions = append(version.Properties.PublishingProfile.TargetRegions, &armcompute.TargetRegion{
-		Name:                 &location,
-		RegionalReplicaCount: to.Ptr[int32](1),
-		StorageAccountType:   to.Ptr(armcompute.StorageAccountTypeStandardLRS),
-	})
+	for _, region := range missing {
+		version.Properties.PublishingProfile.TargetRegions = append(version.Properties.PublishingProfile.TargetRegions, &armcompute.TargetRegion{
+			Name:                 to.Ptr(region),
+			RegionalReplicaCount: to.Ptr[int32](1),
+			StorageAccountType:   to.Ptr(armcompute.StorageAccountTypeStandardLRS),
+		})
+	}
 
 	resp, err := galleryImageVersion.BeginCreateOrUpdate(ctx, image.Gallery.ResourceGroupName, image.Gallery.Name, image.Name, *version.Name, *version, nil)
 	if err != nil {
@@ -806,15 +880,6 @@ func DefaultRetryOpts() policy.RetryOptions {
 	}
 }
 
-func replicatedToCurrentRegion(version *armcompute.GalleryImageVersion, location string) bool {
-	for _, targetRegion := range version.Properties.PublishingProfile.TargetRegions {
-		if strings.EqualFold(strings.ReplaceAll(*targetRegion.Name, " ", ""), location) {
-			return true
-		}
-	}
-	return false
-}
-
 // DeleteSIGImageVersion deletes a SIG image version
 func (a *AzureClient) DeleteSIGImageVersion(ctx context.Context, galleryResourceGroup, galleryName, imageName, version string) {
 	// Ignore errors, don't need to wait for the deletion to complete

e2e/config/regions.go

@@ -0,0 +1,51 @@
+package config
+
+import (
+	"slices"
+	"strings"
+)
+
+// e2eRegions is every region E2E runs in. Scenarios name their region inline; add it here too
+// when a scenario needs a new one, otherwise images are never replicated there and the
+// scenario fails. TestScenarioRegionsAreReplicated catches the common case of forgetting.
+var e2eRegions = []string{
+	"eastus",
+	"southcentralus",
+	"southeastasia",
+	"uaenorth",
+	"westus2",
+	"westus3", // Config.DefaultLocation: where scenarios that name no region run
+}
+
+// replicationRegions returns the regions an image version is replicated to.
+//
+// The result deliberately does not depend on the region the caller wants. TargetRegions is
+// full desired state, so a writer that appends only its own region computes a different
+// desired state from every other writer and its update can drop their regions - which is what
+// made scenarios fail with GalleryImageNotFound. A fixed set makes every writer submit the
+// same list, so a lost update loses nothing and no locking or merging is needed.
+func (i *Image) replicationRegions() []string {
+	switch {
+	case i.Ephemeral:
+		// Throwaway image created and deleted by one test: single writer, and no reason to
+		// pay to replicate it anywhere.
+		return nil
+	case i.OS == OSWindows:
+		// No Windows scenario names a region, and Windows images are large enough that
+		// replicating them where no Windows test looks is pure cost.
+		return []string{"westus3"}
+	default:
+		return e2eRegions
+	}
+}
+
+// SupportsE2ERegion reports whether scenarios using this image are replicated to location.
+func (i *Image) SupportsE2ERegion(location string) bool {
+	return i.Ephemeral || slices.Contains(i.replicationRegions(), NormalizeRegion(location))
+}
+
+// NormalizeRegion converts a region name to the compact lowercase form ARM uses in resource
+// IDs, so "West US 2", "westus2" and "WestUS2" compare equal.
+func NormalizeRegion(location string) string {
+	return strings.ToLower(strings.ReplaceAll(strings.TrimSpace(location), " ", ""))
+}

e2e/config/regions_test.go

@@ -0,0 +1,134 @@
+package config
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+	"os"
+	"path/filepath"
+	"regexp"
+	"slices"
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7"
+)
+
+// A scenario naming a region that images are not replicated to would fail at runtime, so the
+// inline regions are checked against e2eRegions here instead. Regions passed to a helper
+// rather than set on the struct are not visible to this scan; maybeSkipScenario catches those
+// when the scenario runs.
+func TestScenarioRegionsAreReplicated(t *testing.T) {
+	files, err := filepath.Glob("../*_test.go")
+	if err != nil {
+		t.Fatalf("listing scenario files: %v", err)
+	}
+
+	pinned := regexp.MustCompile(`(?m)^\s*Location:\s*"([^"]+)"`)
+	var found int
+	for _, path := range files {
+		source, err := os.ReadFile(path)
+		if err != nil {
+			t.Fatalf("reading %s: %v", path, err)
+		}
+		for _, match := range pinned.FindAllStringSubmatch(string(source), -1) {
+			found++
+			if !slices.Contains(e2eRegions, NormalizeRegion(match[1])) {
+				t.Errorf("%s runs a scenario in %q, which is missing from e2eRegions in config/regions.go", filepath.Base(path), match[1])
+			}
+		}
+	}
+	if found == 0 {
+		t.Fatal("found no scenario regions to check; has the Location field changed shape?")
+	}
+}
+
+// The whole point of the fixed set is that desired state does not depend on the caller. If it
+// ever does again, concurrent writers can clobber each other's regions.
+func TestReplicationRegionsDoNotDependOnCaller(t *testing.T) {
+	image := &Image{OS: OSUbuntu}
+	if got := image.replicationRegions(); !slices.Equal(got, e2eRegions) {
+		t.Fatalf("expected the fixed set %v, got %v", e2eRegions, got)
+	}
+
+	previous := Config.DefaultLocation
+	Config.DefaultLocation = "northeurope"
+	t.Cleanup(func() { Config.DefaultLocation = previous })
+
+	// E2E_LOCATION is per-process, so letting it widen the set would make two pipelines
+	// compute different desired states - the original bug, across processes.
+	if got := image.replicationRegions(); !slices.Equal(got, e2eRegions) {
+		t.Fatalf("E2E_LOCATION must not change the replication set, got %v", got)
+	}
+	if image.SupportsE2ERegion("northeurope") {
+		t.Error("an unlisted region must be unsupported so the run fails fast")
+	}
+}
+
+func TestReplicationRegionsByImageKind(t *testing.T) {
+	if got := (&Image{OS: OSWindows}).replicationRegions(); !slices.Equal(got, []string{"westus3"}) {
+		t.Errorf("Windows images should replicate only to the default location, got %v", got)
+	}
+	if got := (&Image{OS: OSUbuntu, Ephemeral: true}).replicationRegions(); len(got) != 0 {
+		t.Errorf("ephemeral images should opt out of the shared set, got %v", got)
+	}
+	if !(&Image{OS: OSUbuntu, Ephemeral: true}).SupportsE2ERegion("northeurope") {
+		t.Error("ephemeral images should be usable in any region")
+	}
+}
+
+func TestSupportsE2ERegionNormalizesARMFormatting(t *testing.T) {
+	image := &Image{OS: OSUbuntu}
+	for _, location := range []string{"westus2", "West US 2", "WestUS2", " westus2 "} {
+		if !image.SupportsE2ERegion(location) {
+			t.Errorf("expected %q to be recognised as westus2", location)
+		}
+	}
+}
+
+func TestMissingRegions(t *testing.T) {
+	existing := []*armcompute.TargetRegion{
+		{Name: to.Ptr("West US 2")}, // ARM returns display names, not the compact form
+		nil,
+		{Name: nil},
+	}
+	got := missingRegions(existing, []string{"westus2", "eastus"})
+	if !slices.Equal(got, []string{"eastus"}) {
+		t.Fatalf("expected only eastus to be missing, got %v", got)
+	}
+	if got := missingRegions(nil, nil); len(got) != 0 {
+		t.Fatalf("expected no missing regions, got %v", got)
+	}
+}
+
+func TestNormalizeRegion(t *testing.T) {
+	for input, want := range map[string]string{
+		"westus2": "westus2", "West US 2": "westus2", "WESTUS2": "westus2", " westus2 ": "westus2", "": "",
+	} {
+		if got := NormalizeRegion(input); got != want {
+			t.Errorf("NormalizeRegion(%q) = %q, want %q", input, got, want)
+		}
+	}
+}
+
+func TestIsGalleryUpdateConflict(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{"conflict", &azcore.ResponseError{StatusCode: http.StatusConflict}, true},
+		{"precondition failed", &azcore.ResponseError{StatusCode: http.StatusPreconditionFailed}, true},
+		{"wrapped", fmt.Errorf("update: %w", &azcore.ResponseError{StatusCode: http.StatusConflict}), true},
+		{"not found", &azcore.ResponseError{StatusCode: http.StatusNotFound}, false},
+		{"plain", errors.New("boom"), false},
+		{"nil", nil, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := isGalleryUpdateConflict(tc.err); got != tc.want {
+				t.Fatalf("got %v, want %v", got, tc.want)
+			}
+		})
+	}
+}

e2e/config/vhd.go

@@ -321,6 +321,10 @@ type Image struct {
 	IgnoreFailedCgroupTelemetryServices bool
 	Flatcar                             bool
 	SkipOldVHDValidations               bool
+	// Ephemeral marks an image version created at runtime from a single test's disk and
+	// deleted on that test's cleanup. It has exactly one writer, so it is replicated only
+	// where it is used instead of to the shared E2E region set.
+	Ephemeral bool
 	// OSDiskSizeGB overrides the default OS disk size (50 GB) when set.
 	OSDiskSizeGB int32
 }
@@ -401,7 +405,7 @@ func getVHDResourceIDFromMetadata(metadata map[string]vhdMetadataEntry, image Im
 		return "", fmt.Errorf("%w: image %s is not present in E2E VHD metadata", ErrNotFound, image.Name)
 	}
 	for _, region := range entry.Regions {
-		if strings.EqualFold(region, location) {
+		if NormalizeRegion(region) == NormalizeRegion(location) {
 			return entry.ResourceID, nil
 		}
 	}