From 404866ab6d7203eb9be9314a36e4e5e65d32d04f Mon Sep 17 00:00:00 2001
From: Pengfei Ni <peni@microsoft.com>
Date: Fri, 5 Mar 2021 15:21:15 +0800
Subject: [PATCH] Ensure only one LoadBalancer rule is created when HA mode is
 enabled

---
 .../azure/azure_loadbalancer.go               |  7 +++
 .../azure/azure_loadbalancer_test.go          | 53 +++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/staging/src/k8s.io/legacy-cloud-providers/azure/azure_loadbalancer.go b/staging/src/k8s.io/legacy-cloud-providers/azure/azure_loadbalancer.go
index da75a051a7025..68faa939f945e 100644
--- a/staging/src/k8s.io/legacy-cloud-providers/azure/azure_loadbalancer.go
+++ b/staging/src/k8s.io/legacy-cloud-providers/azure/azure_loadbalancer.go
@@ -1630,7 +1630,13 @@ func (az *Cloud) reconcileLoadBalancerRule(
 
 	var expectedProbes []network.Probe
 	var expectedRules []network.LoadBalancingRule
+	highAvailabilityPortsEnabled := false
 	for _, port := range ports {
+		if highAvailabilityPortsEnabled {
+			// Since the port is always 0 when enabling HA, only one rule should be configured.
+			break
+		}
+
 		protocols := []v1.Protocol{port.Protocol}
 		if v, ok := service.Annotations[ServiceAnnotationLoadBalancerMixedProtocols]; ok && v == "true" {
 			klog.V(2).Infof("reconcileLoadBalancerRule lb name (%s) flag(%s) is set", lbName, ServiceAnnotationLoadBalancerMixedProtocols)
@@ -1729,6 +1735,7 @@ func (az *Cloud) reconcileLoadBalancerRule(
 				expectedRule.FrontendPort = to.Int32Ptr(0)
 				expectedRule.BackendPort = to.Int32Ptr(0)
 				expectedRule.Protocol = network.TransportProtocolAll
+				highAvailabilityPortsEnabled = true
 			}
 
 			// we didn't construct the probe objects for UDP or SCTP because they're not allowed on Azure.
diff --git a/staging/src/k8s.io/legacy-cloud-providers/azure/azure_loadbalancer_test.go b/staging/src/k8s.io/legacy-cloud-providers/azure/azure_loadbalancer_test.go
index c8857639f0de6..3504d6f424782 100644
--- a/staging/src/k8s.io/legacy-cloud-providers/azure/azure_loadbalancer_test.go
+++ b/staging/src/k8s.io/legacy-cloud-providers/azure/azure_loadbalancer_test.go
@@ -1592,6 +1592,28 @@ func TestReconcileLoadBalancerRule(t *testing.T) {
 			expectedProbes:  getDefaultTestProbes("http", "/healthy"),
 			expectedRules:   getDefaultTestRules(true),
 		},
+		{
+			desc: "reconcileLoadBalancerRule shall return corresponding probe and lbRule (slb with HA enabled)",
+			service: getTestService("test1", v1.ProtocolTCP, map[string]string{
+				"service.beta.kubernetes.io/azure-load-balancer-enable-high-availability-ports": "true",
+				"service.beta.kubernetes.io/azure-load-balancer-internal":                       "true",
+			}, false, 80),
+			loadBalancerSku: "standard",
+			wantLb:          true,
+			expectedProbes:  getDefaultTestProbes("Tcp", ""),
+			expectedRules:   getHATestRules(true),
+		},
+		{
+			desc: "reconcileLoadBalancerRule shall return corresponding probe and lbRule (slb with HA enabled multi-ports services)",
+			service: getTestService("test1", v1.ProtocolTCP, map[string]string{
+				"service.beta.kubernetes.io/azure-load-balancer-enable-high-availability-ports": "true",
+				"service.beta.kubernetes.io/azure-load-balancer-internal":                       "true",
+			}, false, 80, 8080),
+			loadBalancerSku: "standard",
+			wantLb:          true,
+			expectedProbes:  getDefaultTestProbes("Tcp", ""),
+			expectedRules:   getHATestRules(true),
+		},
 	}
 	for i, test := range testCases {
 		az := GetTestCloud(ctrl)
@@ -1665,6 +1687,37 @@ func getDefaultTestRules(enableTCPReset bool) []network.LoadBalancingRule {
 	return expectedRules
 }
 
+func getHATestRules(enableTCPReset bool) []network.LoadBalancingRule {
+	expectedRules := []network.LoadBalancingRule{
+		{
+			Name: to.StringPtr("atest1-TCP-80"),
+			LoadBalancingRulePropertiesFormat: &network.LoadBalancingRulePropertiesFormat{
+				Protocol: network.TransportProtocol("All"),
+				FrontendIPConfiguration: &network.SubResource{
+					ID: to.StringPtr("frontendIPConfigID"),
+				},
+				BackendAddressPool: &network.SubResource{
+					ID: to.StringPtr("backendPoolID"),
+				},
+				LoadDistribution:     "Default",
+				FrontendPort:         to.Int32Ptr(0),
+				BackendPort:          to.Int32Ptr(0),
+				EnableFloatingIP:     to.BoolPtr(true),
+				DisableOutboundSnat:  to.BoolPtr(false),
+				IdleTimeoutInMinutes: to.Int32Ptr(0),
+				Probe: &network.SubResource{
+					ID: to.StringPtr("/subscriptions/subscription/resourceGroups/rg/providers/" +
+						"Microsoft.Network/loadBalancers/lbname/probes/atest1-TCP-80"),
+				},
+			},
+		},
+	}
+	if enableTCPReset {
+		expectedRules[0].EnableTCPReset = to.BoolPtr(true)
+	}
+	return expectedRules
+}
+
 func getTestLoadBalancer(name, rgName, clusterName, identifier *string, service v1.Service, lbSku string) network.LoadBalancer {
 	lb := network.LoadBalancer{
 		Name: name,