From b6cc31d26684c2d550542b00f9dc362397e9ba13 Mon Sep 17 00:00:00 2001
From: Pengfei Ni <feiskyer@gmail.com>
Date: Tue, 5 Nov 2019 23:41:45 +0800
Subject: [PATCH] Ensure health probes are created for local traffic policy UDP
 services on Azure

---
 .../azure/azure_loadbalancer.go               |  5 +-
 .../azure/azure_loadbalancer_test.go          | 72 +++++++++++++++++++
 2 files changed, 75 insertions(+), 2 deletions(-)

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 3c52863e2a25..001e772e6530 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
@@ -1096,8 +1096,9 @@ func (az *Cloud) reconcileLoadBalancerRule(
 				expectedRule.LoadBalancingRulePropertiesFormat.IdleTimeoutInMinutes = lbIdleTimeout
 			}
 
-			// we didn't construct the probe objects for UDP or SCTP because they're not used/needed/allowed
-			if protocol != v1.ProtocolUDP && protocol != v1.ProtocolSCTP {
+			// we didn't construct the probe objects for UDP or SCTP because they're not allowed on Azure.
+			// However, when externalTrafficPolicy is Local, Kubernetes HTTP health check would be used for probing.
+			if servicehelpers.NeedsHealthCheck(service) || (protocol != v1.ProtocolUDP && protocol != v1.ProtocolSCTP) {
 				expectedRule.Probe = &network.SubResource{
 					ID: to.StringPtr(az.getLoadBalancerProbeID(lbName, lbRuleName)),
 				}
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 808b35790108..5a5bf4fe59b5 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
@@ -1330,6 +1330,60 @@ func TestReconcileLoadBalancer(t *testing.T) {
 		},
 	}
 
+	service6 := getTestService("test1", v1.ProtocolUDP, nil, 80)
+	lb6 := getTestLoadBalancer(to.StringPtr("testCluster"), to.StringPtr("testCluster"), to.StringPtr("atest1"), service6, "basic")
+	lb6.FrontendIPConfigurations = &[]network.FrontendIPConfiguration{}
+	lb6.Probes = &[]network.Probe{}
+	expectedLB6 := getTestLoadBalancer(to.StringPtr("testCluster"), to.StringPtr("testCluster"), to.StringPtr("atest1"), service6, "basic")
+	expectedLB6.Probes = &[]network.Probe{}
+	(*expectedLB6.LoadBalancerPropertiesFormat.LoadBalancingRules)[0].Probe = nil
+	(*expectedLB6.LoadBalancerPropertiesFormat.LoadBalancingRules)[0].EnableTCPReset = nil
+	(*expectedLB6.LoadBalancerPropertiesFormat.LoadBalancingRules)[0].DisableOutboundSnat = to.BoolPtr(false)
+	expectedLB6.FrontendIPConfigurations = &[]network.FrontendIPConfiguration{
+		{
+			Name: to.StringPtr("atest1"),
+			FrontendIPConfigurationPropertiesFormat: &network.FrontendIPConfigurationPropertiesFormat{
+				PublicIPAddress: &network.PublicIPAddress{ID: to.StringPtr("/subscriptions/subscription/" +
+					"resourceGroups/rg/providers/Microsoft.Network/publicIPAddresses/pipName")},
+			},
+		},
+	}
+
+	service7 := getTestService("test1", v1.ProtocolUDP, nil, 80)
+	service7.Spec.HealthCheckNodePort = 10081
+	service7.Spec.ExternalTrafficPolicy = v1.ServiceExternalTrafficPolicyTypeLocal
+	lb7 := getTestLoadBalancer(to.StringPtr("testCluster"), to.StringPtr("testCluster"), to.StringPtr("atest1"), service7, "basic")
+	lb7.FrontendIPConfigurations = &[]network.FrontendIPConfiguration{}
+	lb7.Probes = &[]network.Probe{}
+	expectedLB7 := getTestLoadBalancer(to.StringPtr("testCluster"), to.StringPtr("testCluster"), to.StringPtr("atest1"), service7, "basic")
+	(*expectedLB7.LoadBalancerPropertiesFormat.LoadBalancingRules)[0].Probe = &network.SubResource{
+		ID: to.StringPtr("/subscriptions/subscription/resourceGroups/rg/providers/Microsoft.Network/loadBalancers/testCluster/probes/atest1-UDP-80"),
+	}
+	(*expectedLB7.LoadBalancerPropertiesFormat.LoadBalancingRules)[0].EnableTCPReset = nil
+	(*expectedLB7.LoadBalancerPropertiesFormat.LoadBalancingRules)[0].DisableOutboundSnat = to.BoolPtr(false)
+	expectedLB7.FrontendIPConfigurations = &[]network.FrontendIPConfiguration{
+		{
+			Name: to.StringPtr("atest1"),
+			FrontendIPConfigurationPropertiesFormat: &network.FrontendIPConfigurationPropertiesFormat{
+				PublicIPAddress: &network.PublicIPAddress{ID: to.StringPtr("/subscriptions/subscription/" +
+					"resourceGroups/rg/providers/Microsoft.Network/publicIPAddresses/pipName")},
+			},
+		},
+	}
+	expectedLB7.Probes = &[]network.Probe{
+		{
+			Name: to.StringPtr("atest1-" + string(service7.Spec.Ports[0].Protocol) +
+				"-" + strconv.Itoa(int(service7.Spec.Ports[0].Port))),
+			ProbePropertiesFormat: &network.ProbePropertiesFormat{
+				Port:              to.Int32Ptr(10081),
+				RequestPath:       to.StringPtr("/healthz"),
+				Protocol:          network.ProbeProtocolHTTP,
+				IntervalInSeconds: to.Int32Ptr(5),
+				NumberOfProbes:    to.Int32Ptr(2),
+			},
+		},
+	}
+
 	testCases := []struct {
 		desc                string
 		service             v1.Service
@@ -1389,6 +1443,24 @@ func TestReconcileLoadBalancer(t *testing.T) {
 			expectedLB:          expectedSLb5,
 			expectedError:       nil,
 		},
+		{
+			desc:            "reconcileLoadBalancer shall reconcile UDP services",
+			loadBalancerSku: "basic",
+			service:         service6,
+			existingLB:      lb6,
+			wantLb:          true,
+			expectedLB:      expectedLB6,
+			expectedError:   nil,
+		},
+		{
+			desc:            "reconcileLoadBalancer shall reconcile probes for local traffic policy UDP services",
+			loadBalancerSku: "basic",
+			service:         service7,
+			existingLB:      lb7,
+			wantLb:          true,
+			expectedLB:      expectedLB7,
+			expectedError:   nil,
+		},
 	}
 
 	for i, test := range testCases {