From f59ad10c388b05358dba8f4cc86872d7b4bb4cdf 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      | 27 +++++++++++++++++++++++++-
 azure/services/agentpools/spec_test.go | 16 +++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/azure/services/agentpools/spec.go b/azure/services/agentpools/spec.go
index 14c3f035e83..efc89e23a5f 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,18 +261,22 @@ 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 == "" {
 			// agent pool is up to date, nothing to do
 			log.V(4).Info("no changes found between user-updated spec and existing spec")
+			fmt.Printf("no changes found between user-updated spec and existing spec\n")
 			return nil, nil
 		}
 		log.V(4).Info("found a diff between the desired spec and the existing agentpool", "difference", diff)
+		fmt.Printf("found a diff between the desired spec and the existing agentpool: %s\n", diff)
 	}
 
 	availabilityZones := azure.PtrSlice(&s.AvailabilityZones)
-	nodeTaints := azure.PtrSlice(&s.NodeTaints)
 	var sku *string
 	if s.SKU != "" {
 		sku = &s.SKU
@@ -402,3 +407,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..a9a6fa11411 100644
--- a/azure/services/agentpools/spec_test.go
+++ b/azure/services/agentpools/spec_test.go
@@ -357,6 +357,22 @@ func TestParameters(t *testing.T) {
 			expected:      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("fake-taint")}
+				},
+				sdkWithProvisioningState("Succeeded"),
+			),
+			expected:      nil,
+			expectedError: nil,
+		},
 		{
 			name: "parameters with an existing agent pool and update needed on node taints",
 			spec: fakeAgentPool(),