From 9c92c882df5e0fd9fc8e4fd60fdacfbf00f4b023 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      | 37 +++++++++++++++++++++---
 azure/services/agentpools/spec_test.go | 39 ++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/azure/services/agentpools/spec.go b/azure/services/agentpools/spec.go
index 2d6af10c9c5..1aa85b7ce21 100644
--- a/azure/services/agentpools/spec.go
+++ b/azure/services/agentpools/spec.go
@@ -168,6 +168,10 @@ func (s *AgentPoolSpec) Parameters(ctx context.Context, existing interface{}) (p
 	defer done()
 
 	nodeLabels := s.NodeLabels
+	var nodeTaints *[]string
+	if len(s.NodeTaints) > 0 {
+		nodeTaints = &s.NodeTaints
+	}
 	if existing != nil {
 		existingPool, ok := existing.(containerservice.AgentPool)
 		if !ok {
@@ -243,6 +247,9 @@ func (s *AgentPoolSpec) Parameters(ctx context.Context, existing interface{}) (p
 		nodeLabels = mergeSystemNodeLabels(normalizedProfile.NodeLabels, existingPool.NodeLabels)
 		normalizedProfile.NodeLabels = nodeLabels
 
+		nodeTaints = mergeSystemNodeTaints(normalizedProfile.NodeTaints, existingPool.NodeTaints)
+		normalizedProfile.NodeTaints = nodeTaints
+
 		// Compute a diff to check if we require an update
 		diff := cmp.Diff(normalizedProfile, existingProfile)
 		if diff == "" {
@@ -257,10 +264,6 @@ func (s *AgentPoolSpec) Parameters(ctx context.Context, existing interface{}) (p
 	if len(s.AvailabilityZones) > 0 {
 		availabilityZones = &s.AvailabilityZones
 	}
-	var nodeTaints *[]string
-	if len(s.NodeTaints) > 0 {
-		nodeTaints = &s.NodeTaints
-	}
 	var sku *string
 	if s.SKU != "" {
 		sku = &s.SKU
@@ -384,3 +387,29 @@ 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 {
+	var ret []string
+	if capz != nil {
+		ret = *capz
+	}
+	if ret == nil {
+		ret = make([]string, 0)
+	}
+	// Look for labels returned from the AKS node pool API that begin with kubernetes.azure.com
+	if aks != nil {
+		for _, aksNodeTaint := range *aks {
+			if azureutil.IsAzureSystemNodeLabelKey(aksNodeTaint) {
+				ret = append(ret, aksNodeTaint)
+			}
+		}
+	}
+
+	// Preserve nil-ness of capz
+	if capz == nil && len(ret) == 0 {
+		return nil
+	}
+	return &ret
+}
diff --git a/azure/services/agentpools/spec_test.go b/azure/services/agentpools/spec_test.go
index a64f0998ab9..a4271ecc892 100644
--- a/azure/services/agentpools/spec_test.go
+++ b/azure/services/agentpools/spec_test.go
@@ -29,6 +29,7 @@ import (
 	"k8s.io/utils/pointer"
 	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 {
@@ -126,6 +127,12 @@ func sdkWithProvisioningState(state string) func(*containerservice.AgentPool) {
 	}
 }
 
+func sdkWithNodeTaints(nodeTaints *[]string) func(*containerservice.AgentPool) {
+	return func(pool *containerservice.AgentPool) {
+		pool.ManagedClusterAgentPoolProfileProperties.NodeTaints = nodeTaints
+	}
+}
+
 func TestParameters(t *testing.T) {
 	testcases := []struct {
 		name          string
@@ -299,6 +306,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 *containerservice.AgentPool) {
+					pool.NodeTaints = &[]string{"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 *containerservice.AgentPool) {
+					pool.NodeTaints = &[]string{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(),