From 9552b25050bff4f2b9d912485bef1c3e71e933f5 Mon Sep 17 00:00:00 2001
From: Peter Rifel <pgrifel@gmail.com>
Date: Thu, 8 Jul 2021 22:41:00 -0400
Subject: [PATCH] Azure - support VMSS availability zones

Azure's subnets are regional so we use similar functionality to GCE where we reference the InstanceGroup's zones rather than a subnet's zone.
IG Zones are already populated on cluster creation here: https://github.com/kubernetes/kops/blob/b35803789627345f58eb979cf53fa87bfe422f28/upup/pkg/fi/cloudup/new_cluster.go#L682-L684
---
 pkg/model/azuremodel/vmscaleset.go                | 1 +
 upup/pkg/fi/cloudup/azuretasks/vmscaleset.go      | 7 ++++++-
 upup/pkg/fi/cloudup/azuretasks/vmscaleset_test.go | 9 +++++++++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/pkg/model/azuremodel/vmscaleset.go b/pkg/model/azuremodel/vmscaleset.go
index 9cd716ca7c..7f1cd6c35f 100644
--- a/pkg/model/azuremodel/vmscaleset.go
+++ b/pkg/model/azuremodel/vmscaleset.go
@@ -78,6 +78,7 @@ func (b *VMScaleSetModelBuilder) buildVMScaleSetTask(
 		SKUName:            fi.String(ig.Spec.MachineType),
 		ComputerNamePrefix: fi.String(ig.Name),
 		AdminUser:          fi.String(b.Cluster.Spec.CloudConfig.Azure.AdminUser),
+		Zones:              ig.Spec.Zones,
 	}
 
 	var err error
diff --git a/upup/pkg/fi/cloudup/azuretasks/vmscaleset.go b/upup/pkg/fi/cloudup/azuretasks/vmscaleset.go
index 402dc0252b..e84ae13bae 100644
--- a/upup/pkg/fi/cloudup/azuretasks/vmscaleset.go
+++ b/upup/pkg/fi/cloudup/azuretasks/vmscaleset.go
@@ -116,6 +116,7 @@ type VMScaleSet struct {
 	// CustomData is the user data configuration
 	CustomData  fi.Resource
 	Tags        map[string]*string
+	Zones       []string
 	PrincipalID *string
 }
 
@@ -231,6 +232,9 @@ func (s *VMScaleSet) Find(c *fi.Context) (*VMScaleSet, error) {
 			Name: to.StringPtr(loadBalancerID.LoadBalancerName),
 		}
 	}
+	if found.Zones != nil {
+		vmss.Zones = *found.Zones
+	}
 	return vmss, nil
 }
 
@@ -367,7 +371,8 @@ func (s *VMScaleSet) RenderAzure(t *azure.AzureAPITarget, a, e, changes *VMScale
 		Identity: &compute.VirtualMachineScaleSetIdentity{
 			Type: compute.ResourceIdentityTypeSystemAssigned,
 		},
-		Tags: e.Tags,
+		Tags:  e.Tags,
+		Zones: &e.Zones,
 	}
 
 	result, err := t.Cloud.VMScaleSet().CreateOrUpdate(
diff --git a/upup/pkg/fi/cloudup/azuretasks/vmscaleset_test.go b/upup/pkg/fi/cloudup/azuretasks/vmscaleset_test.go
index 9859eb7ab2..2df118fdec 100644
--- a/upup/pkg/fi/cloudup/azuretasks/vmscaleset_test.go
+++ b/upup/pkg/fi/cloudup/azuretasks/vmscaleset_test.go
@@ -86,6 +86,7 @@ func newTestVMScaleSet() *VMScaleSet {
 		SSHPublicKey:       to.StringPtr("ssh"),
 		CustomData:         fi.NewStringResource("custom"),
 		Tags:               map[string]*string{},
+		Zones:              []string{"zone1"},
 	}
 }
 
@@ -128,6 +129,10 @@ func TestVMScaleSetRenderAzure(t *testing.T) {
 	if expected.PrincipalID == nil {
 		t.Errorf("unexpected nil principalID")
 	}
+
+	if a, e := *actual.Zones, expected.Zones; !reflect.DeepEqual(a, e) {
+		t.Errorf("unexpected Zone: expected %s, but got %s", e, a)
+	}
 }
 
 func TestVMScaleSetFind(t *testing.T) {
@@ -253,6 +258,7 @@ func TestVMScaleSetFind(t *testing.T) {
 		Identity: &compute.VirtualMachineScaleSetIdentity{
 			Type: compute.ResourceIdentityTypeSystemAssigned,
 		},
+		Zones: &[]string{"zone1"},
 	}
 	if _, err := cloud.VMScaleSet().CreateOrUpdate(context.Background(), *rg.Name, *vmss.Name, vmssParameters); err != nil {
 		t.Fatalf("failed to create: %s", err)
@@ -288,6 +294,9 @@ func TestVMScaleSetFind(t *testing.T) {
 	if !*actual.RequirePublicIP {
 		t.Errorf("unexpected require public IP")
 	}
+	if a, e := actual.Zones, *vmssParameters.Zones; !reflect.DeepEqual(a, e) {
+		t.Errorf("unexpected Zone: expected %s, but got %s", e, a)
+	}
 }
 
 func TestVMScaleSetRun(t *testing.T) {