From 01acc1c7a4a02313d79f205b0931fd2141596d85 Mon Sep 17 00:00:00 2001
From: Alex Leites <18728999+tallaxes@users.noreply.github.com>
Date: Tue, 28 Nov 2023 14:51:10 -0800
Subject: [PATCH] fix: respect NodePool zone requirement (if any) (#47)

* fix: respect zone requirements

* test: zone-aware provisioning
---
 pkg/providers/instance/instance.go       |  7 +++++--
 pkg/providers/instancetype/suite_test.go | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/pkg/providers/instance/instance.go b/pkg/providers/instance/instance.go
index 3f7c9afb0..54ade243f 100644
--- a/pkg/providers/instance/instance.go
+++ b/pkg/providers/instance/instance.go
@@ -553,8 +553,11 @@ func (p *Provider) pickSkuSizePriorityAndZone(ctx context.Context, nodeClaim *co
 	logging.FromContext(ctx).Infof("Selected instance type %s", instanceType.Name)
 	// Priority - Provisioner defaults to Regular, so pick Spot if it is explicitly included in requirements (and is offered in at least one zone)
 	priority := p.getPriorityForInstanceType(nodeClaim, instanceType)
-	// Zone - ideally random/spread from zones that support given Priority
-	priorityOfferings := lo.Filter(instanceType.Offerings.Available(), func(o corecloudprovider.Offering, _ int) bool { return o.CapacityType == priority })
+	// Zone - ideally random/spread from requested zones that support given Priority
+	requestedZones := scheduling.NewNodeSelectorRequirements(nodeClaim.Spec.Requirements...).Get(v1.LabelTopologyZone)
+	priorityOfferings := lo.Filter(instanceType.Offerings.Available(), func(o corecloudprovider.Offering, _ int) bool {
+		return o.CapacityType == priority && requestedZones.Has(o.Zone)
+	})
 	zonesWithPriority := lo.Map(priorityOfferings, func(o corecloudprovider.Offering, _ int) string { return o.Zone })
 	if zone, ok := sets.New(zonesWithPriority...).PopAny(); ok {
 		// Zones in Offerings have <region>-<number> format; the zone returned from here will be used for VM instantiation,
diff --git a/pkg/providers/instancetype/suite_test.go b/pkg/providers/instancetype/suite_test.go
index c51d765f6..e5b3416c6 100644
--- a/pkg/providers/instancetype/suite_test.go
+++ b/pkg/providers/instancetype/suite_test.go
@@ -880,6 +880,25 @@ var _ = Describe("InstanceType Provider", func() {
 		})
 	})
 
+	Context("Zone aware provisioning", func() {
+		It("should launch in the NodePool-requested zone", func() {
+			zone, vmZone := "eastus-3", "3"
+			nodePool.Spec.Template.Spec.Requirements = []v1.NodeSelectorRequirement{
+				{Key: corev1beta1.CapacityTypeLabelKey, Operator: v1.NodeSelectorOpIn, Values: []string{corev1beta1.CapacityTypeSpot, corev1beta1.CapacityTypeOnDemand}},
+				{Key: v1.LabelTopologyZone, Operator: v1.NodeSelectorOpIn, Values: []string{zone}},
+			}
+			ExpectApplied(ctx, env.Client, nodePool, nodeClass)
+			pod := coretest.UnschedulablePod()
+			ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, coreProvisioner, pod)
+			node := ExpectScheduled(ctx, env.Client, pod)
+			Expect(node.Labels).To(HaveKeyWithValue(v1alpha2.AlternativeLabelTopologyZone, zone))
+
+			vm := azureEnv.VirtualMachinesAPI.VirtualMachineCreateOrUpdateBehavior.CalledWithInput.Pop().VM
+			Expect(vm).NotTo(BeNil())
+			Expect(vm.Zones).To(ConsistOf(&vmZone))
+		})
+	})
+
 })
 
 var _ = Describe("Tax Calculator", func() {