fix(e2e): replicate gallery images to a fixed region set so concurrent writers converge

e2e/config/azure.go

@@ -652,30 +652,237 @@ func (a *AzureClient) LatestSIGImageVersionByTag(ctx context.Context, image *Ima
 	return VHDResourceID(*latestVersion.ID), nil
 }
 
+// ensureReplication makes an image version usable in location. For shared images it
+// replicates to the whole E2E region set rather than just location: TargetRegions is full
+// desired state, so writers that each add only their own region can clobber one another,
+// while writers that all submit the same superset converge no matter how their reads and
+// writes interleave.
 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)
+	start := time.Now()
+
+	if err := a.ensureTargetRegions(ctx, image, version, image.replicationRegions(location)); err != nil {
+		// The extra regions only prevent future contention; this test needs one region.
+		// Fail only if that region did not make it, so a quota, a region disabled on the
+		// subscription, or contention over regions this test never touches cannot fail it.
+		live, getErr := a.getImageVersion(ctx, image, *version.Name)
+		if getErr != nil || !hasTargetRegion(live.Properties.PublishingProfile.TargetRegions, location) {
+			return err
+		}
+		toolkit.Logf(ctx, "Replicating %s to the full E2E region set failed, but %s is a target so continuing: %v", *version.ID, location, err)
+		*version = *live
 	}
 
-	if replicatedToCurrentRegion(version, location) {
-		toolkit.Logf(ctx, "Image version %s is already in region %s", *version.ID, location)
-		return nil
+	// Being listed in TargetRegions only records intent. The regional replica can still be
+	// "Replicating" while the parent ProvisioningState is already "Succeeded" — that gap is
+	// the source of GalleryImageNotFound 404s on VMSS create, so wait for the region itself.
+	if err := a.waitForRegionalReplicationCompleted(ctx, image, version, location); err != nil {
+		return err
 	}
-	regions := make([]string, 0, len(version.Properties.PublishingProfile.TargetRegions))
-	for _, targetRegion := range version.Properties.PublishingProfile.TargetRegions {
-		regions = append(regions, *targetRegion.Name)
+
+	elapsed := time.Since(start)
+	toolkit.LogDuration(ctx, elapsed, 3*time.Minute, fmt.Sprintf("Replication took: %s (%s)", elapsed, *version.ID))
+	return nil
+}
+
+// ensureTargetRegions adds every missing region in a single update. On conflict it re-reads
+// and merges again, because a concurrent writer may have updated the version in between.
+func (a *AzureClient) ensureTargetRegions(ctx context.Context, image *Image, version *armcompute.GalleryImageVersion, desiredRegions []string) error {
+	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)
+		}
+
+		targetRegions, missing := mergeTargetRegions(version.Properties.PublishingProfile.TargetRegions, desiredRegions)
+		if len(missing) == 0 {
+			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, ", "))
+
+		previous := version.Properties.PublishingProfile.TargetRegions
+		version.Properties.PublishingProfile.TargetRegions = targetRegions
+		err := a.updateImageVersion(ctx, image, version)
+		if err == nil {
+			return nil
+		}
+		version.Properties.PublishingProfile.TargetRegions = previous
+		if !isGalleryUpdateConflict(err) || attempt >= maxAttempts {
+			return err
+		}
+
+		toolkit.Logf(ctx, "Concurrent update of image version %s; re-reading and merging again", *version.ID)
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(5 * time.Second):
+		}
+
+		live, err := a.getImageVersion(ctx, image, *version.Name)
+		if err != nil {
+			return err
+		}
+		*version = *live
 	}
-	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)
+}
 
-	start := time.Now() // Record the start time
-	err := a.replicateImageVersionToCurrentRegion(ctx, image, version, location)
-	elapsed := time.Since(start) // Calculate the elapsed time
+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)
+	}
+	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)
+	}
+	return &resp.GalleryImageVersion, nil
+}
 
-	toolkit.LogDuration(ctx, elapsed, 3*time.Minute, fmt.Sprintf("Replication took: %s (%s)", elapsed, *version.ID))
+func (a *AzureClient) updateImageVersion(ctx context.Context, image *Image, version *armcompute.GalleryImageVersion) error {
+	client, err := armcompute.NewGalleryImageVersionsClient(image.Gallery.SubscriptionID, a.Credential, a.ArmOptions)
+	if err != nil {
+		return fmt.Errorf("create a new images client: %v", err)
+	}
+	op, err := client.BeginCreateOrUpdate(ctx, image.Gallery.ResourceGroupName, image.Gallery.Name, image.Name, *version.Name, *version, nil)
+	if err != nil {
+		return fmt.Errorf("begin updating image version target regions: %w", err)
+	}
+	if _, err := op.PollUntilDone(ctx, DefaultPollUntilDoneOptions); err != nil {
+		return fmt.Errorf("updating image version target regions: %w", err)
+	}
+	return nil
+}
+
+// mergeTargetRegions returns the target regions with any missing desired region added, along
+// with the names that were missing. Existing entries are preserved as-is so per-region replica
+// counts and storage account types configured elsewhere are never rewritten.
+func mergeTargetRegions(existing []*armcompute.TargetRegion, desiredRegions []string) ([]*armcompute.TargetRegion, []string) {
+	merged := append([]*armcompute.TargetRegion(nil), existing...)
+	known := make(map[string]struct{}, len(existing))
+	for _, region := range existing {
+		if region != nil && region.Name != nil {
+			known[NormalizeRegion(*region.Name)] = struct{}{}
+		}
+	}
+
+	var missing []string
+	for _, location := range desiredRegions {
+		normalized := NormalizeRegion(location)
+		if normalized == "" {
+			continue
+		}
+		if _, exists := known[normalized]; exists {
+			continue
+		}
+		known[normalized] = struct{}{}
+		missing = append(missing, normalized)
+		merged = append(merged, &armcompute.TargetRegion{
+			Name:                 to.Ptr(normalized),
+			RegionalReplicaCount: to.Ptr[int32](1),
+			StorageAccountType:   to.Ptr(armcompute.StorageAccountTypeStandardLRS),
+		})
+	}
+	slices.Sort(missing)
+	return merged, missing
+}
 
-	return err
+func isGalleryUpdateConflict(err error) bool {
+	var respErr *azcore.ResponseError
+	return errors.As(err, &respErr) &&
+		(respErr.StatusCode == http.StatusConflict || respErr.StatusCode == http.StatusPreconditionFailed)
+}
+
+// waitForRegionalReplicationCompleted polls the gallery image version with the
+// ReplicationStatus expand option until the named region's RegionalReplicationStatus.State
+// is Completed. Returns an error if the regional replication enters a terminal Failed state
+// or the wait times out. Treats Unknown as transient (Azure occasionally emits Unknown
+// briefly during state transitions).
+//
+// This closes a documented gap in the SIG API: the parent provisioning LRO can succeed
+// before per-region replicas are actually serving traffic. Without this wait, callers see
+// GalleryImageNotFound 404s on VMSS creation that look like bugs but are infra eventual
+// consistency.
+func (a *AzureClient) waitForRegionalReplicationCompleted(ctx context.Context, image *Image, version *armcompute.GalleryImageVersion, location string) error {
+	imgVersionClient, err := armcompute.NewGalleryImageVersionsClient(image.Gallery.SubscriptionID, a.Credential, a.ArmOptions)
+	if err != nil {
+		return fmt.Errorf("create image version client for replication wait: %w", err)
+	}
+
+	getOpts := &armcompute.GalleryImageVersionsClientGetOptions{
+		Expand: to.Ptr(armcompute.ReplicationStatusTypesReplicationStatus),
+	}
+
+	var lastLoggedState armcompute.ReplicationState
+	const (
+		pollInterval = 15 * time.Second
+		pollTimeout  = 20 * time.Minute
+	)
+
+	pollErr := wait.PollUntilContextTimeout(ctx, pollInterval, pollTimeout, true, func(ctx context.Context) (bool, error) {
+		resp, err := imgVersionClient.Get(ctx, image.Gallery.ResourceGroupName, image.Gallery.Name, image.Name, *version.Name, getOpts)
+		if err != nil {
+			// Transient API errors should not abort the wait; the SDK already retries throttling.
+			toolkit.Logf(ctx, "transient error getting image version replication status (will retry): %v", err)
+			return false, nil
+		}
+
+		regional := findRegionalReplicationStatus(resp.GalleryImageVersion.Properties.ReplicationStatus, location)
+		if regional == nil || regional.State == nil {
+			// Region not yet present in the status summary; keep waiting.
+			return false, nil
+		}
+
+		state := *regional.State
+		if state != lastLoggedState {
+			progress := int32(0)
+			if regional.Progress != nil {
+				progress = *regional.Progress
+			}
+			toolkit.Logf(ctx, "Image version %s regional replication state in %s: %s (progress: %d%%)", *version.ID, location, state, progress)
+			lastLoggedState = state
+		}
+
+		switch state {
+		case armcompute.ReplicationStateCompleted:
+			return true, nil
+		case armcompute.ReplicationStateFailed:
+			details := ""
+			if regional.Details != nil {
+				details = *regional.Details
+			}
+			return false, fmt.Errorf("regional replication of %s to %s failed: %s", *version.ID, location, details)
+		case armcompute.ReplicationStateReplicating, armcompute.ReplicationStateUnknown:
+			return false, nil
+		default:
+			// Forward-compat: unknown future states treated as in-progress.
+			return false, nil
+		}
+	})
+
+	if pollErr != nil {
+		return fmt.Errorf("waiting for regional replication of %s to %s to complete: %w", *version.ID, location, pollErr)
+	}
+	toolkit.Logf(ctx, "Image version %s regional replication to %s confirmed Completed", *version.ID, location)
+	return nil
+}
+
+// findRegionalReplicationStatus finds the RegionalReplicationStatus for the given location
+// (case-insensitive, ignoring spaces) within the version-wide ReplicationStatus.
+func findRegionalReplicationStatus(status *armcompute.ReplicationStatus, location string) *armcompute.RegionalReplicationStatus {
+	if status == nil {
+		return nil
+	}
+	normalized := NormalizeRegion(location)
+	for _, regional := range status.Summary {
+		if regional == nil || regional.Region == nil {
+			continue
+		}
+		if NormalizeRegion(*regional.Region) == normalized {
+			return regional
+		}
+	}
+	return nil
 }
 
 func (a *AzureClient) waitForVersionOperationCompletion(ctx context.Context, image *Image, version *armcompute.GalleryImageVersion) error {
@@ -732,28 +939,6 @@ 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 {
-	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),
-	})
-
-	resp, err := galleryImageVersion.BeginCreateOrUpdate(ctx, image.Gallery.ResourceGroupName, image.Gallery.Name, image.Name, *version.Name, *version, nil)
-	if err != nil {
-		return fmt.Errorf("begin updating image version target regions: %w", err)
-	}
-	if _, err := resp.PollUntilDone(ctx, DefaultPollUntilDoneOptions); err != nil {
-		return fmt.Errorf("updating image version target regions: %w", err)
-	}
-
-	return nil
-}
-
 func (a *AzureClient) EnsureSIGImageVersion(ctx context.Context, image *Image, location string) (VHDResourceID, error) {
 	galleryImageVersion, err := armcompute.NewGalleryImageVersionsClient(image.Gallery.SubscriptionID, a.Credential, a.ArmOptions)
 	if err != nil {
@@ -783,6 +968,39 @@ func (a *AzureClient) EnsureSIGImageVersion(ctx context.Context, image *Image, l
 	return VHDResourceID(*resp.ID), nil
 }
 
+// WaitForImageVersionReplicatedToRegion is the public entrypoint for callers (e.g. the
+// VMSS create retry loop) that need to confirm a previously-resolved image is actually
+// serving traffic in a region before retrying a request that just failed with
+// GalleryImageNotFound. It fetches the live image version with ReplicationStatus
+// expanded and blocks until the regional state is Completed (or fails fast on terminal
+// Failed). It does not initiate replication; callers should have already gone through
+// EnsureSIGImageVersion / LatestSIGImageVersionByTag.
+func (a *AzureClient) WaitForImageVersionReplicatedToRegion(ctx context.Context, image *Image, location string) error {
+	if image == nil {
+		return fmt.Errorf("nil image")
+	}
+	if image.Version == "" {
+		return fmt.Errorf("image %s has no resolved version; cannot check replication state", image.Name)
+	}
+	live, err := a.getImageVersion(ctx, image, image.Version)
+	if err != nil {
+		return fmt.Errorf("get image version %s/%s for replication-wait: %w", image.Name, image.Version, err)
+	}
+	return a.waitForRegionalReplicationCompleted(ctx, image, live, location)
+}
+
+// hasTargetRegion reports whether location is already a replication target, comparing region
+// names case- and space-insensitively because ARM returns both "West US 2" and "westus2".
+func hasTargetRegion(regions []*armcompute.TargetRegion, location string) bool {
+	normalized := NormalizeRegion(location)
+	for _, region := range regions {
+		if region != nil && region.Name != nil && NormalizeRegion(*region.Name) == normalized {
+			return true
+		}
+	}
+	return false
+}
+
 func DefaultRetryOpts() policy.RetryOptions {
 	return policy.RetryOptions{
 		// Use generous retry settings to survive Azure Compute Gallery throttling.
@@ -806,15 +1024,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