[ARO-25489] Reduce the memory usage of ResourceSKUs querying

pkg/cluster/loadbalancerinternal.go

@@ -60,7 +60,7 @@ func MigrateInternalLoadBalancerZones(
 		}
 	}
 
-	filteredSkus, err := computeskus.GetVMSkusForCurrentRegion(ctx, resourceSkusClient, location)
+	filteredSkus, err := computeskus.SelectVMSkusInCurrentRegion(ctx, resourceSkusClient, location, []string{string(doc.OpenShiftCluster.Properties.MasterProfile.VMSize)})
 	if err != nil {
 		return doc, err
 	}

pkg/cluster/loadbalancerinternal_test.go

@@ -94,19 +94,20 @@ func TestUpdateLoadBalancerZonalNoopAndErrorPaths(t *testing.T) {
 					}, nil,
 				)
 
-				sku.EXPECT().List(gomock.Any(), "location eq eastus", false).Return([]*armcompute.ResourceSKU{
-					{
-						Name:      pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)),
-						Locations: pointerutils.ToSlicePtr([]string{"eastus"}),
-						LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{
-							{
-								Zones: []*string{},
-							},
-						}),
-						Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}),
-						ResourceType: pointerutils.ToPtr("virtualMachines"),
-					},
-				}, nil)
+				sku.EXPECT().List(gomock.Any(), "location eq eastus", false).
+					Return(func(yield func(*armcompute.ResourceSKU, error) bool) {
+						yield(&armcompute.ResourceSKU{
+							Name:      pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)),
+							Locations: pointerutils.ToSlicePtr([]string{"eastus"}),
+							LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{
+								{
+									Zones: []*string{},
+								},
+							}),
+							Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}),
+							ResourceType: pointerutils.ToPtr("virtualMachines"),
+						}, nil)
+					})
 			},
 			expectedLogs: []testlog.ExpectedLogEntry{
 				{
@@ -136,15 +137,8 @@ func TestUpdateLoadBalancerZonalNoopAndErrorPaths(t *testing.T) {
 					}, nil,
 				)
 
-				sku.EXPECT().List(gomock.Any(), "location eq eastus", false).Return([]*armcompute.ResourceSKU{
-					{
-						Name:         pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)),
-						Locations:    pointerutils.ToSlicePtr([]string{"eastus"}),
-						LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{}),
-						Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}),
-						ResourceType: pointerutils.ToPtr("virtualMachines"),
-					},
-				}, nil)
+				sku.EXPECT().List(gomock.Any(), "location eq eastus", false).
+					Return(func(yield func(*armcompute.ResourceSKU, error) bool) {})
 			},
 			expectedLogs: []testlog.ExpectedLogEntry{},
 			wantErrs:     []error{errVMAvailability},
@@ -170,7 +164,10 @@ func TestUpdateLoadBalancerZonalNoopAndErrorPaths(t *testing.T) {
 					}, nil,
 				)
 
-				sku.EXPECT().List(gomock.Any(), "location eq eastus", false).Return([]*armcompute.ResourceSKU{}, errTestSKUFetchError)
+				sku.EXPECT().List(gomock.Any(), "location eq eastus", false).
+					Return(func(yield func(*armcompute.ResourceSKU, error) bool) {
+						yield(nil, errTestSKUFetchError)
+					})
 			},
 			expectedLogs: []testlog.ExpectedLogEntry{},
 			wantErrs:     []error{computeskus.ErrListVMResourceSKUs, errTestSKUFetchError},
@@ -380,17 +377,18 @@ func TestUpdateLoadBalancerZonalMigration(t *testing.T) {
 				}, nil,
 			)
 
-			skus.EXPECT().List(gomock.Any(), "location eq eastus", false).Return([]*armcompute.ResourceSKU{
-				{
-					Name:      pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)),
-					Locations: pointerutils.ToSlicePtr([]string{"eastus"}),
-					LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{
-						{Zones: pointerutils.ToSlicePtr([]string{"1", "2", "3"})},
-					}),
-					Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}),
-					ResourceType: pointerutils.ToPtr("virtualMachines"),
-				},
-			}, nil)
+			skus.EXPECT().List(gomock.Any(), "location eq eastus", false).
+				Return(func(yield func(*armcompute.ResourceSKU, error) bool) {
+					yield(&armcompute.ResourceSKU{
+						Name:      pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)),
+						Locations: pointerutils.ToSlicePtr([]string{"eastus"}),
+						LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{
+							{Zones: pointerutils.ToSlicePtr([]string{"1", "2", "3"})},
+						}),
+						Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}),
+						ResourceType: pointerutils.ToPtr("virtualMachines"),
+					}, nil)
+				})
 
 			plsFIPRemoval := plses.EXPECT().CreateOrUpdateAndWait(gomock.Any(), rgName, infraID+"-pls", armnetwork.PrivateLinkService{
 				Properties: &armnetwork.PrivateLinkServiceProperties{

pkg/cluster/validate.go

@@ -32,7 +32,11 @@ func (m *manager) validateResources(ctx context.Context) error {
 // created. This function is only to be called during cluster bootstrap!
 func (m *manager) validateZones(ctx context.Context) error {
 	location := m.doc.OpenShiftCluster.Location
-	filteredSkus, err := computeskus.GetVMSkusForCurrentRegion(ctx, m.armResourceSKUs, location)
+
+	filteredSkus, err := computeskus.SelectVMSkusInCurrentRegion(ctx, m.armResourceSKUs, location, []string{
+		string(m.doc.OpenShiftCluster.Properties.MasterProfile.VMSize),
+		string(m.doc.OpenShiftCluster.Properties.WorkerProfiles[0].VMSize),
+	})
 	if err != nil {
 		return err
 	}

pkg/cluster/validate_test.go

@@ -6,7 +6,6 @@ package cluster
 import (
 	"context"
 	"errors"
-	"fmt"
 	"strings"
 	"testing"
 
@@ -180,8 +179,18 @@ func TestValidateZones(t *testing.T) {
 
 			resourceSkusClient := mock_armcompute.NewMockResourceSKUsClient(controller)
 			resourceSkusClient.EXPECT().
-				List(gomock.Any(), fmt.Sprintf("location eq %v", "eastus"), false).
-				Return(skus, tt.resourceSkusClientErr)
+				List(gomock.Any(), "location eq eastus", false).
+				Return(func(yield func(*armcompute.ResourceSKU, error) bool) {
+					if tt.resourceSkusClientErr != nil {
+						yield(nil, tt.resourceSkusClientErr)
+						return
+					}
+					for _, v := range skus {
+						if !yield(v, nil) {
+							return
+						}
+					}
+				})
 
 			m := &manager{
 				env: mockEnv,

pkg/frontend/admin_openshiftcluster_vmresize_pre_validation.go

@@ -28,7 +28,6 @@ import (
 	arov1alpha1 "github.com/Azure/ARO-RP/pkg/operator/apis/aro.openshift.io/v1alpha1"
 	"github.com/Azure/ARO-RP/pkg/util/azureclient/mgmt/compute"
 	"github.com/Azure/ARO-RP/pkg/util/clusteroperators"
-	"github.com/Azure/ARO-RP/pkg/util/computeskus"
 )
 
 // getPreResizeControlPlaneVMsValidation is the HTTP handler; the underscore
@@ -337,15 +336,12 @@ func (f *frontend) validateVMSKU(
 		return err
 	}
 
-	skus, err := a.VMSizeList(ctx)
+	filteredSkus, err := a.VMGetSKUs(ctx, []string{desiredVMSize})
 	if err != nil {
 		return err
 	}
 
 	location := doc.OpenShiftCluster.Location
-
-	filteredSkus := computeskus.FilterVMSizes(skus, location)
-
 	sku, err := checkSKUAvailability(filteredSkus, location, "vmSize", desiredVMSize)
 	if err != nil {
 		return err

pkg/frontend/admin_openshiftcluster_vmresize_pre_validation_test.go

@@ -142,9 +142,9 @@ func TestPreResizeControlPlaneVMsValidation(t *testing.T) {
 			},
 			mocks: func(tt *test, a *mock_adminactions.MockAzureActions) {
 				a.EXPECT().
-					VMSizeList(gomock.Any()).
-					Return([]*armcompute.ResourceSKU{
-						{
+					VMGetSKUs(gomock.Any(), []string{"Standard_D8s_v3"}).
+					Return(map[string]*armcompute.ResourceSKU{
+						"Standard_D8s_v3": {
 							Name:         pointerutils.ToPtr("Standard_D8s_v3"),
 							ResourceType: pointerutils.ToPtr("virtualMachines"),
 							Locations:    pointerutils.ToSlicePtr([]string{"eastus"}),
@@ -308,21 +308,8 @@ func TestPreResizeControlPlaneVMsValidation(t *testing.T) {
 			},
 			mocks: func(tt *test, a *mock_adminactions.MockAzureActions) {
 				a.EXPECT().
-					VMSizeList(gomock.Any()).
-					Return([]*armcompute.ResourceSKU{
-						{
-							Name:         pointerutils.ToPtr("Standard_D16s_v3"),
-							ResourceType: pointerutils.ToPtr("virtualMachines"),
-							Locations:    pointerutils.ToSlicePtr([]string{"eastus"}),
-							LocationInfo: []*armcompute.ResourceSKULocationInfo{
-								{
-									Location: pointerutils.ToPtr("eastus"),
-								},
-							},
-							Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}),
-							Capabilities: []*armcompute.ResourceSKUCapabilities{},
-						},
-					}, nil)
+					VMGetSKUs(gomock.Any(), []string{"Standard_D8s_v3"}).
+					Return(map[string]*armcompute.ResourceSKU{}, nil)
 			},
 			kubeMocks:      allKubeChecksHealthyMock,
 			wantStatusCode: http.StatusBadRequest,
@@ -360,9 +347,9 @@ func TestPreResizeControlPlaneVMsValidation(t *testing.T) {
 			},
 			mocks: func(tt *test, a *mock_adminactions.MockAzureActions) {
 				a.EXPECT().
-					VMSizeList(gomock.Any()).
-					Return([]*armcompute.ResourceSKU{
-						{
+					VMGetSKUs(gomock.Any(), []string{"Standard_D8s_v3"}).
+					Return(map[string]*armcompute.ResourceSKU{
+						"Standard_D8s_v3": {
 							Name:         pointerutils.ToPtr("Standard_D8s_v3"),
 							ResourceType: pointerutils.ToPtr("virtualMachines"),
 							Locations:    pointerutils.ToSlicePtr([]string{"eastus"}),

pkg/frontend/admin_openshiftcluster_vmsizelist.go

@@ -9,13 +9,12 @@ import (
 	"fmt"
 	"net/http"
 	"path/filepath"
+	"slices"
 	"strings"
 
 	"github.com/go-chi/chi/v5"
 	"github.com/sirupsen/logrus"
 
-	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7"
-
 	"github.com/Azure/ARO-RP/pkg/api"
 	"github.com/Azure/ARO-RP/pkg/database/cosmosdb"
 	"github.com/Azure/ARO-RP/pkg/frontend/middleware"
@@ -66,19 +65,8 @@ func (f *frontend) _getAdminOpenShiftClusterVMResizeOptions(ctx context.Context,
 		return nil, err
 	}
 
-	return json.Marshal(f.filterVMSkus(skus))
-}
-
-func (f *frontend) filterVMSkus(skus []*armcompute.ResourceSKU) []string {
-	filteredSkus := []string{}
-
-	for _, sku := range skus {
-		if sku.Restrictions != nil && len(sku.Restrictions) == 0 {
-			if sku.Name != nil {
-				filteredSkus = append(filteredSkus, *sku.Name)
-			}
-		}
-	}
+	// Sort for stability
+	slices.Sort(skus)
 
-	return filteredSkus
+	return json.Marshal(skus)
 }

pkg/frontend/admin_openshiftcluster_vmsizelist_test.go

@@ -13,14 +13,11 @@ import (
 	"github.com/sirupsen/logrus"
 	"go.uber.org/mock/gomock"
 
-	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7"
-
 	"github.com/Azure/ARO-RP/pkg/api"
 	"github.com/Azure/ARO-RP/pkg/env"
 	"github.com/Azure/ARO-RP/pkg/frontend/adminactions"
 	"github.com/Azure/ARO-RP/pkg/metrics/noop"
 	mock_adminactions "github.com/Azure/ARO-RP/pkg/util/mocks/adminactions"
-	"github.com/Azure/ARO-RP/pkg/util/pointerutils"
 	testdatabase "github.com/Azure/ARO-RP/test/database"
 )
 
@@ -68,23 +65,10 @@ func TestAdminListVMSizeList(t *testing.T) {
 			mocks: func(tt *test, a *mock_adminactions.MockAzureActions) {
 				a.EXPECT().
 					VMSizeList(gomock.Any()).
-					Return([]*armcompute.ResourceSKU{
-						{
-							Name: pointerutils.ToPtr("Standard_D8s_v90"),
-							Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{
-								{
-									Type: pointerutils.ToPtr(armcompute.ResourceSKURestrictionsTypeLocation),
-								},
-							}),
-						},
-						{
-							Name:         pointerutils.ToPtr("Standard_D8s_v9001"),
-							Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}),
-						},
-					}, nil)
+					Return([]string{"Standard_D8s_v9001", "Standard_D8s_v80"}, nil)
 			},
 			wantStatusCode: http.StatusOK,
-			wantResponse:   []byte(`["Standard_D8s_v9001"]` + "\n"),
+			wantResponse:   []byte(`["Standard_D8s_v80","Standard_D8s_v9001"]` + "\n"),
 		},
 		{
 			name:       "cluster not found",

pkg/frontend/adminactions/azureactions.go

@@ -22,6 +22,7 @@ import (
 	"github.com/Azure/ARO-RP/pkg/util/azureclient/mgmt/compute"
 	"github.com/Azure/ARO-RP/pkg/util/azureclient/mgmt/features"
 	"github.com/Azure/ARO-RP/pkg/util/azureclient/mgmt/storage"
+	"github.com/Azure/ARO-RP/pkg/util/computeskus"
 	"github.com/Azure/ARO-RP/pkg/util/stringutils"
 )
 
@@ -34,7 +35,8 @@ type AzureActions interface {
 	VMRedeployAndWait(ctx context.Context, vmName string) error
 	VMStartAndWait(ctx context.Context, vmName string) error
 	VMStopAndWait(ctx context.Context, vmName string, deallocateVM bool) error
-	VMSizeList(ctx context.Context) ([]*sdkcompute.ResourceSKU, error)
+	VMSizeList(ctx context.Context) ([]string, error)
+	VMGetSKUs(ctx context.Context, vmSizes []string) (map[string]*sdkcompute.ResourceSKU, error)
 	VMResize(ctx context.Context, vmName string, vmSize string) error
 	ResourceGroupHasVM(ctx context.Context, vmName string) (bool, error)
 	VMSerialConsole(ctx context.Context, log *logrus.Entry, vmName string, target io.Writer) error
@@ -148,9 +150,12 @@ func (a *azureActions) VMStopAndWait(ctx context.Context, vmName string, dealloc
 	return a.virtualMachines.StopAndWait(ctx, clusterRGName, vmName, deallocateVM)
 }
 
-func (a *azureActions) VMSizeList(ctx context.Context) ([]*sdkcompute.ResourceSKU, error) {
-	filter := fmt.Sprintf("location eq '%s'", a.env.Location())
-	return a.resourceSkus.List(ctx, filter, false)
+func (a *azureActions) VMGetSKUs(ctx context.Context, vmSizes []string) (map[string]*sdkcompute.ResourceSKU, error) {
+	return computeskus.SelectVMSkusInCurrentRegion(ctx, a.resourceSkus, a.env.Location(), vmSizes)
+}
+
+func (a *azureActions) VMSizeList(ctx context.Context) ([]string, error) {
+	return computeskus.ListUnrestrictedVMSkusInCurrentRegion(ctx, a.resourceSkus, a.env.Location())
 }
 
 func (a *azureActions) VMResize(ctx context.Context, vmName string, size string) error {

pkg/frontend/sku_test.go

@@ -95,7 +95,7 @@ func TestValidateVMSku(t *testing.T) {
 			workerProfile1Sku:      "Standard_D4s_v2",
 			workerProfile2Sku:      "Standard_D4s_v2",
 			resourceSkusClientErr:  errors.New("unable to retrieve skus information"),
-			wantErr:                "unable to retrieve skus information",
+			wantErr:                "failure listing resource SKUs: unable to retrieve skus information",
 			masterEncryptionAtHost: api.EncryptionAtHostDisabled,
 			workerEncryptionAtHost: api.EncryptionAtHostDisabled,
 		},
@@ -320,7 +320,17 @@ func TestValidateVMSku(t *testing.T) {
 			resourceSkusClient := mock_armcompute.NewMockResourceSKUsClient(controller)
 			resourceSkusClient.EXPECT().
 				List(gomock.Any(), fmt.Sprintf("location eq %v", "eastus"), false).
-				Return(skus, tt.resourceSkusClientErr)
+				Return(func(yield func(*armcompute.ResourceSKU, error) bool) {
+					if tt.resourceSkusClientErr != nil {
+						yield(nil, tt.resourceSkusClientErr)
+						return
+					}
+					for _, v := range skus {
+						if !yield(v, nil) {
+							return
+						}
+					}
+				})
 
 			err := validateVMSku(context.Background(), oc, resourceSkusClient)
 			utilerror.AssertErrorMessage(t, err, tt.wantErr)

pkg/frontend/sku_validation.go

@@ -42,14 +42,18 @@ func validateVMSku(ctx context.Context, oc *api.OpenShiftCluster, resourceSkusCl
 	// Get a list of available worker SKUs, filtering by location. We initialized a new resourceSkusClient
 	// so that we can determine SKU availability within target cluster subscription instead of within RP subscription.
 	location := oc.Location
-	filter := fmt.Sprintf("location eq %s", location)
-	skus, err := resourceSkusClient.List(ctx, filter, false)
+
+	skus := []string{string(oc.Properties.MasterProfile.VMSize)}
+	workerProfiles, _ := api.GetEnrichedWorkerProfiles(oc.Properties)
+	for _, workerprofile := range workerProfiles {
+		skus = append(skus, string(workerprofile.VMSize))
+	}
+
+	filteredSkus, err := computeskus.SelectVMSkusInCurrentRegion(ctx, resourceSkusClient, location, skus)
 	if err != nil {
 		return err
 	}
 
-	filteredSkus := computeskus.FilterVMSizes(skus, location)
-
 	controlPlaneSKU, err := checkSKUAvailability(filteredSkus, location, "properties.masterProfile.VMSize", string(oc.Properties.MasterProfile.VMSize))
 	if err != nil {
 		return err
@@ -67,8 +71,6 @@ func validateVMSku(ctx context.Context, oc *api.OpenShiftCluster, resourceSkusCl
 		}
 	}
 
-	workerProfiles, _ := api.GetEnrichedWorkerProfiles(oc.Properties)
-
 	// In case there are multiple WorkerProfiles listed in the cluster document (such as post-install),
 	// compare VMSize in each WorkerProfile to the resourceSkusClient call above to ensure that the sku is available in region.
 	// XXX: Will this ever be called post-install?