From fc44a4e0a5275e0f794218def09d457baa4c0afa Mon Sep 17 00:00:00 2001
From: willie-yao <williamyao@microsoft.com>
Date: Mon, 2 Oct 2023 23:35:01 +0000
Subject: [PATCH] Fix AKS reconciliation of taints

---
 azure/services/agentpools/spec.go      | 25 ++++++++++++++++-
 azure/services/agentpools/spec_test.go | 39 ++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/azure/services/agentpools/spec.go b/azure/services/agentpools/spec.go
index 14c3f035e83..9d75a36d38f 100644
--- a/azure/services/agentpools/spec.go
+++ b/azure/services/agentpools/spec.go
@@ -178,6 +178,7 @@ func (s *AgentPoolSpec) Parameters(ctx context.Context, existing interface{}) (p
 	defer done()
 
 	nodeLabels := s.NodeLabels
+	nodeTaints := azure.PtrSlice(&s.NodeTaints)
 	if existing != nil {
 		existingPool, ok := existing.(armcontainerservice.AgentPool)
 		if !ok {
@@ -260,6 +261,9 @@ func (s *AgentPoolSpec) Parameters(ctx context.Context, existing interface{}) (p
 		nodeLabels = mergeSystemNodeLabels(normalizedProfile.Properties.NodeLabels, existingPool.Properties.NodeLabels)
 		normalizedProfile.Properties.NodeLabels = nodeLabels
 
+		nodeTaints = mergeSystemNodeTaints(normalizedProfile.Properties.NodeTaints, existingPool.Properties.NodeTaints)
+		normalizedProfile.Properties.NodeTaints = nodeTaints
+
 		// Compute a diff to check if we require an update
 		diff := cmp.Diff(normalizedProfile, existingProfile)
 		if diff == "" {
@@ -271,7 +275,6 @@ func (s *AgentPoolSpec) Parameters(ctx context.Context, existing interface{}) (p
 	}
 
 	availabilityZones := azure.PtrSlice(&s.AvailabilityZones)
-	nodeTaints := azure.PtrSlice(&s.NodeTaints)
 	var sku *string
 	if s.SKU != "" {
 		sku = &s.SKU
@@ -402,3 +405,23 @@ func mergeSystemNodeLabels(capz, aks map[string]*string) map[string]*string {
 	}
 	return ret
 }
+
+// mergeSystemNodeTaints appends any kubernetes.azure.com-prefixed taints from the AKS taint set
+// into the local capz taint set.
+func mergeSystemNodeTaints(capz, aks []*string) []*string {
+	ret := capz
+	if ret == nil {
+		ret = make([]*string, 0)
+	}
+	// Look for taints returned from the AKS node pool API that begin with kubernetes.azure.com
+	for _, aksNodeTaint := range aks {
+		if azureutil.IsAzureSystemNodeLabelKey(*aksNodeTaint) {
+			ret = append(ret, aksNodeTaint)
+		}
+	}
+	// Preserve nil-ness of capz
+	if capz == nil && len(ret) == 0 {
+		ret = nil
+	}
+	return ret
+}
diff --git a/azure/services/agentpools/spec_test.go b/azure/services/agentpools/spec_test.go
index bf9cdf0cb34..48b64d238c7 100644
--- a/azure/services/agentpools/spec_test.go
+++ b/azure/services/agentpools/spec_test.go
@@ -30,6 +30,7 @@ import (
 	"k8s.io/utils/ptr"
 	infrav1 "sigs.k8s.io/cluster-api-provider-azure/api/v1beta1"
 	"sigs.k8s.io/cluster-api-provider-azure/azure"
+	azureutil "sigs.k8s.io/cluster-api-provider-azure/util/azure"
 )
 
 func fakeAgentPool(changes ...func(*AgentPoolSpec)) AgentPoolSpec {
@@ -149,6 +150,12 @@ func sdkWithSpotMaxPrice(spotMaxPrice float32) func(*armcontainerservice.AgentPo
 	}
 }
 
+func sdkWithNodeTaints(nodeTaints []*string) func(*armcontainerservice.AgentPool) {
+	return func(pool *armcontainerservice.AgentPool) {
+		pool.Properties.NodeTaints = nodeTaints
+	}
+}
+
 func TestParameters(t *testing.T) {
 	testcases := []struct {
 		name          string
@@ -357,6 +364,38 @@ func TestParameters(t *testing.T) {
 			expected:      nil,
 			expectedError: nil,
 		},
+		{
+			name: "difference in non-system node taints with empty taints should trigger update",
+			spec: fakeAgentPool(
+				func(pool *AgentPoolSpec) {
+					pool.NodeTaints = nil
+				},
+			),
+			existing: sdkFakeAgentPool(
+				func(pool *armcontainerservice.AgentPool) {
+					pool.Properties.NodeTaints = []*string{ptr.To("fake-taint")}
+				},
+				sdkWithProvisioningState("Succeeded"),
+			),
+			expected:      sdkFakeAgentPool(sdkWithNodeTaints(nil)),
+			expectedError: nil,
+		},
+		{
+			name: "difference in system node taints with empty taints shouldn't trigger update",
+			spec: fakeAgentPool(
+				func(pool *AgentPoolSpec) {
+					pool.NodeTaints = nil
+				},
+			),
+			existing: sdkFakeAgentPool(
+				func(pool *armcontainerservice.AgentPool) {
+					pool.Properties.NodeTaints = []*string{ptr.To(azureutil.AzureSystemNodeLabelPrefix + "-fake-taint")}
+				},
+				sdkWithProvisioningState("Succeeded"),
+			),
+			expected:      nil,
+			expectedError: nil,
+		},
 		{
 			name: "parameters with an existing agent pool and update needed on node taints",
 			spec: fakeAgentPool(),