azure: Distribute compute subnets to proper zones.

Distributing the compute subnets across NAT gateways. Depends
on nat gateway availability zones and vm zones.

Author: rna-afk

pkg/asset/installconfig/azure/metadata.go

@@ -6,6 +6,8 @@ import (
 	"sort"
 	"sync"
 
+	"sigs.k8s.io/cluster-api-provider-azure/api/v1beta1"
+
 	typesazure "github.com/openshift/installer/pkg/types/azure"
 )
 
@@ -17,7 +19,9 @@ type Metadata struct {
 	client            API
 	dnsCfg            *DNSConfig
 	availabilityZones []string
+	vmZones           []string
 	region            string
+	ZonesSubnetMap    map[string][]string
 
 	// CloudName indicates the Azure cloud environment (e.g. public, gov't).
 	CloudName typesazure.CloudEnvironment `json:"cloudName,omitempty"`
@@ -122,3 +126,75 @@ func (m *Metadata) AvailabilityZones(ctx context.Context) ([]string, error) {
 
 	return m.availabilityZones, nil
 }
+
+// VMAvailabilityZones retrieves a list of availability zones for the configured region and instance type.
+func (m *Metadata) VMAvailabilityZones(ctx context.Context, instanceType string) ([]string, error) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	if len(m.vmZones) == 0 {
+		zones, err := m.client.GetAvailabilityZones(ctx, m.region, instanceType)
+		if err != nil {
+			return nil, fmt.Errorf("error retrieving Availability Zones: %w", err)
+		}
+		if zones != nil {
+			sort.Strings(zones)
+			m.vmZones = zones
+		}
+	}
+
+	return m.vmZones, nil
+}
+
+// GenerateZonesSubnetMap creates a map of all the zones that are supported for nat gateways and vms and
+// sets it to the subnets provided. If no subnets are provided, it creates subnets for multi zone
+// functionality.
+func (m *Metadata) GenerateZonesSubnetMap(subnetSpec []typesazure.SubnetSpec, defaultComputeSubnet string) (map[string][]string, error) {
+	if m.ZonesSubnetMap == nil {
+		// Get the availability zones.
+		if m.availabilityZones == nil {
+			_, err := m.AvailabilityZones(context.TODO())
+			if err != nil {
+				return nil, err
+			}
+		}
+		subnetZones := m.availabilityZones
+		computeSubnets := []string{}
+
+		// Get all the byo subnets or generate subnet per az.
+		if len(subnetSpec) != 0 {
+			sort.Slice(subnetSpec, func(i, j int) bool {
+				return subnetSpec[i].Name < subnetSpec[j].Name
+			})
+			for _, subnet := range subnetSpec {
+				if subnet.Role == v1beta1.SubnetNode {
+					computeSubnets = append(computeSubnets, subnet.Name)
+				}
+			}
+		} else {
+			for idx := range subnetZones {
+				computeName := fmt.Sprintf("%s-%d", defaultComputeSubnet, idx)
+				if idx == 0 {
+					computeName = defaultComputeSubnet
+				}
+				computeSubnets = append(computeSubnets, computeName)
+			}
+		}
+
+		// Assign zone to subnets.
+		subnetMap := map[string][]string{}
+		zoneIndex := 0
+		for _, subnet := range computeSubnets {
+			if _, ok := subnetMap[subnetZones[zoneIndex]]; !ok {
+				subnetMap[subnetZones[zoneIndex]] = []string{}
+			}
+			subnetMap[subnetZones[zoneIndex]] = append(subnetMap[subnetZones[zoneIndex]], subnet)
+			zoneIndex++
+			if zoneIndex >= len(subnetZones) {
+				zoneIndex = 0
+			}
+		}
+		m.ZonesSubnetMap = subnetMap
+	}
+	return m.ZonesSubnetMap, nil
+}

pkg/asset/machines/azure/machinesets.go

@@ -2,20 +2,24 @@ package azure
 
 import (
 	"fmt"
+	"slices"
 	"sort"
 
 	"github.com/pkg/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/util/sets"
 
 	clusterapi "github.com/openshift/api/machine/v1beta1"
+	"github.com/openshift/installer/pkg/asset/installconfig"
 	icazure "github.com/openshift/installer/pkg/asset/installconfig/azure"
 	"github.com/openshift/installer/pkg/types"
 	"github.com/openshift/installer/pkg/types/azure"
 )
 
 // MachineSets returns a list of machinesets for a machinepool.
-func MachineSets(clusterID string, config *types.InstallConfig, pool *types.MachinePool, osImage, role, userDataSecret string, capabilities map[string]string, useImageGallery bool, subnetZones []string, session *icazure.Session) ([]*clusterapi.MachineSet, error) {
+func MachineSets(clusterID string, ic *installconfig.InstallConfig, pool *types.MachinePool, osImage, role, userDataSecret string, capabilities map[string]string, useImageGallery bool, subnetZones []string, session *icazure.Session) ([]*clusterapi.MachineSet, error) {
+	config := ic.Config
 	if configPlatform := config.Platform.Name(); configPlatform != azure.Name {
 		return nil, fmt.Errorf("non-azure configuration: %q", configPlatform)
 	}
@@ -47,6 +51,27 @@ func MachineSets(clusterID string, config *types.InstallConfig, pool *types.Mach
 	sort.Strings(azs)
 	subnetIndex := -1
 	var machinesets []*clusterapi.MachineSet
+
+	if config.Azure.OutboundType == azure.NATGatewayMultiZoneOutboundType {
+		return getMultiZoneMachineSets(multiZoneMachineSetInput{
+			networkResourceGroup: networkResourceGroup,
+			virtualNetworkName:   virtualNetworkName,
+			platform:             platform,
+			mpool:                mpool,
+			osImage:              osImage,
+			userDataSecret:       userDataSecret,
+			clusterID:            clusterID,
+			role:                 role,
+			capabilities:         capabilities,
+			useImageGallery:      useImageGallery,
+			session:              session,
+			subnetSpec:           config.Azure.Subnets,
+			replicas:             total,
+			ic:                   ic,
+			azs:                  azs,
+			pool:                 pool,
+		})
+	}
 	for idx, az := range azs {
 		replicas := int32(total / numOfAZs)
 		if int64(idx) < total%numOfAZs {
@@ -102,3 +127,115 @@ func MachineSets(clusterID string, config *types.InstallConfig, pool *types.Mach
 	}
 	return machinesets, nil
 }
+
+type multiZoneMachineSetInput struct {
+	networkResourceGroup string
+	platform             *azure.Platform
+	mpool                *azure.MachinePool
+	osImage              string
+	userDataSecret       string
+	clusterID            string
+	role                 string
+	capabilities         map[string]string
+	useImageGallery      bool
+	session              *icazure.Session
+	virtualNetworkName   string
+	subnetSpec           []azure.SubnetSpec
+	replicas             int64
+	ic                   *installconfig.InstallConfig
+	azs                  []string
+	pool                 *types.MachinePool
+}
+
+func getMultiZoneMachineSets(in multiZoneMachineSetInput) ([]*clusterapi.MachineSet, error) {
+	// Deep copy metadata map.
+	zoneSubnetmap := map[string][]string{}
+	subnetCount := 0
+	for key, value := range in.ic.Azure.ZonesSubnetMap {
+		if slices.Contains(in.azs, key) {
+			zoneSubnetmap[key] = sets.NewString(value...).List()
+			subnetCount += len(value)
+		}
+	}
+	machineSets := []*clusterapi.MachineSet{}
+	replicasToCreate := int32(in.replicas)
+	replicaPerSet := max(replicasToCreate/int32(subnetCount), 1)
+	remainder := replicasToCreate % int32(subnetCount)
+	if replicasToCreate < int32(subnetCount) {
+		remainder = 0
+	}
+	numAZUsed := map[string]int{}
+	for _, az := range in.azs {
+		numAZUsed[az] = 0
+	}
+	for replicasToCreate != 0 && len(zoneSubnetmap) != 0 {
+		for idx, az := range in.azs {
+			if _, ok := zoneSubnetmap[az]; !ok {
+				continue
+			}
+			subnet := zoneSubnetmap[az][0]
+			if len(zoneSubnetmap[az]) == 1 {
+				delete(zoneSubnetmap, az)
+			} else {
+				zoneSubnetmap[az] = zoneSubnetmap[az][1:]
+			}
+			currentReplica := replicaPerSet
+			if remainder != 0 {
+				currentReplica++
+				remainder--
+			}
+			provider, err := provider(in.platform, in.mpool, in.osImage, in.userDataSecret, in.clusterID, in.role, &idx, in.capabilities, in.useImageGallery, in.session, in.networkResourceGroup, in.virtualNetworkName, subnet)
+			if err != nil {
+				return nil, errors.Wrap(err, "failed to create provider")
+			}
+			name := fmt.Sprintf("%s-%s-%s%s-%d", in.clusterID, in.pool.Name, in.platform.Region, az, numAZUsed[az])
+			if numAZUsed[az] == 0 {
+				name = fmt.Sprintf("%s-%s-%s%s", in.clusterID, in.pool.Name, in.platform.Region, az)
+			}
+			numAZUsed[az]++
+			mset := &clusterapi.MachineSet{
+				TypeMeta: metav1.TypeMeta{
+					APIVersion: "machine.openshift.io/v1beta1",
+					Kind:       "MachineSet",
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "openshift-machine-api",
+					Name:      name,
+					Labels: map[string]string{
+						"machine.openshift.io/cluster-api-cluster":      in.clusterID,
+						"machine.openshift.io/cluster-api-machine-role": in.role,
+						"machine.openshift.io/cluster-api-machine-type": in.role,
+					},
+				},
+				Spec: clusterapi.MachineSetSpec{
+					Replicas: &currentReplica,
+					Selector: metav1.LabelSelector{
+						MatchLabels: map[string]string{
+							"machine.openshift.io/cluster-api-machineset": name,
+							"machine.openshift.io/cluster-api-cluster":    in.clusterID,
+						},
+					},
+					Template: clusterapi.MachineTemplateSpec{
+						ObjectMeta: clusterapi.ObjectMeta{
+							Labels: map[string]string{
+								"machine.openshift.io/cluster-api-machineset":   name,
+								"machine.openshift.io/cluster-api-cluster":      in.clusterID,
+								"machine.openshift.io/cluster-api-machine-role": in.role,
+								"machine.openshift.io/cluster-api-machine-type": in.role,
+							},
+						},
+						Spec: clusterapi.MachineSpec{
+							ProviderSpec: clusterapi.ProviderSpec{
+								Value: &runtime.RawExtension{Object: provider},
+							},
+							// we don't need to set Versions, because we control those via cluster operators.
+						},
+					},
+				},
+			}
+			machineSets = append(machineSets, mset)
+			replicasToCreate -= currentReplica
+		}
+	}
+	return machineSets, nil
+}

pkg/asset/machines/clusterapi.go

@@ -236,7 +236,7 @@ func (c *ClusterAPI) Generate(ctx context.Context, dependencies asset.Parents) e
 		}
 
 		if len(mpool.Zones) == 0 {
-			azs, err := client.GetAvailabilityZones(ctx, ic.Platform.Azure.Region, mpool.InstanceType)
+			azs, err := installConfig.Azure.VMAvailabilityZones(ctx, mpool.InstanceType)
 			if err != nil {
 				return fmt.Errorf("failed to fetch availability zones: %w", err)
 			}

pkg/asset/machines/master.go

@@ -370,7 +370,7 @@ func (m *Master) Generate(ctx context.Context, dependencies asset.Parents) error
 		}
 
 		if len(mpool.Zones) == 0 {
-			azs, err := client.GetAvailabilityZones(ctx, ic.Platform.Azure.Region, mpool.InstanceType)
+			azs, err := installConfig.Azure.VMAvailabilityZones(ctx, mpool.InstanceType)
 			if err != nil {
 				return errors.Wrap(err, "failed to fetch availability zones")
 			}

pkg/asset/machines/worker.go

@@ -571,7 +571,7 @@ func (w *Worker) Generate(ctx context.Context, dependencies asset.Parents) error
 			}
 
 			if len(mpool.Zones) == 0 {
-				azs, err := client.GetAvailabilityZones(ctx, ic.Platform.Azure.Region, mpool.InstanceType)
+				azs, err := installConfig.Azure.VMAvailabilityZones(ctx, mpool.InstanceType)
 				if err != nil {
 					return errors.Wrap(err, "failed to fetch availability zones")
 				}
@@ -588,6 +588,11 @@ func (w *Worker) Generate(ctx context.Context, dependencies asset.Parents) error
 				if err != nil {
 					return errors.Wrap(err, "failed to fetch availability zones")
 				}
+				computeSubnet := installConfig.Config.Azure.ComputeSubnetName(clusterID.InfraID)
+				_, err := installConfig.Azure.GenerateZonesSubnetMap(installConfig.Config.Azure.Subnets, computeSubnet)
+				if err != nil {
+					return err
+				}
 			}
 
 			if mpool.OSImage.Publisher != "" {
@@ -610,7 +615,7 @@ func (w *Worker) Generate(ctx context.Context, dependencies asset.Parents) error
 			}
 
 			useImageGallery := ic.Platform.Azure.CloudName != azuretypes.StackCloud
-			sets, err := azure.MachineSets(clusterID.InfraID, ic, &pool, rhcosImage.Compute, "worker", workerUserDataSecretName, capabilities, useImageGallery, subnetZones, session)
+			sets, err := azure.MachineSets(clusterID.InfraID, installConfig, &pool, rhcosImage.Compute, "worker", workerUserDataSecretName, capabilities, useImageGallery, subnetZones, session)
 			if err != nil {
 				return errors.Wrap(err, "failed to create worker machine objects")
 			}