From daed4c809679e268f31bfa51b0ff323192d9ca14 Mon Sep 17 00:00:00 2001
From: Mike Beaumont <mjboamail@gmail.com>
Date: Wed, 6 Nov 2024 10:12:25 +0100
Subject: [PATCH] fix(MeshLoadBalancingStrategy): set all priorities equal if
 localityAware is disabled (#11980)

We already set `priority: 1` for remote zones in
`fillRemoteMeshServices` so we need to explicitly set `priority: 0` if
we _don't_ want locality awareness.

Also fix an insufficient condition in an existing case, where we weren't
verifying that requests really go to all zones, only that it's OK if
they do.

Fixes #11968

Signed-off-by: Mike Beaumont <mjboamail@gmail.com>
---
 .../plugin/v1alpha1/plugin.go                 |  9 +++++
 .../localityawarelb/meshmultizoneservice.go   | 38 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/pkg/plugins/policies/meshloadbalancingstrategy/plugin/v1alpha1/plugin.go b/pkg/plugins/policies/meshloadbalancingstrategy/plugin/v1alpha1/plugin.go
index 55288733f415..99ab9e3c8719 100644
--- a/pkg/plugins/policies/meshloadbalancingstrategy/plugin/v1alpha1/plugin.go
+++ b/pkg/plugins/policies/meshloadbalancingstrategy/plugin/v1alpha1/plugin.go
@@ -233,6 +233,15 @@ func configureEndpoints(
 			}
 		}
 	}
+	if conf.LocalityAwareness != nil && pointer.Deref(conf.LocalityAwareness.Disabled) {
+		for _, cla := range endpoints {
+			for _, localityLbEndpoints := range cla.Endpoints {
+				if localityLbEndpoints.Locality != nil && localityLbEndpoints.Locality.Zone != localZone {
+					localityLbEndpoints.Priority = 0
+				}
+			}
+		}
+	}
 	return nil
 }
 
diff --git a/test/e2e_env/multizone/localityawarelb/meshmultizoneservice.go b/test/e2e_env/multizone/localityawarelb/meshmultizoneservice.go
index ccdf7fe16e30..78b8948520c6 100644
--- a/test/e2e_env/multizone/localityawarelb/meshmultizoneservice.go
+++ b/test/e2e_env/multizone/localityawarelb/meshmultizoneservice.go
@@ -49,6 +49,7 @@ spec:
 				testserver.WithMesh(meshName),
 				testserver.WithEchoArgs("echo", "--instance", "kube-test-server-1"),
 			)).
+			Install(democlient.Install(democlient.WithNamespace(namespace), democlient.WithMesh(meshName))).
 			SetupInGroup(multizone.KubeZone1, &group)
 
 		NewClusterSetup().
@@ -95,7 +96,9 @@ spec:
 	It("should fallback only to first zone", func() {
 		// given traffic to other zones
 		Eventually(responseFromInstance(multizone.KubeZone2), "30s", "1s").
-			MustPassRepeatedly(5).Should(Or(Equal("kube-test-server-1"), Equal("uni-test-server")))
+			Should(Equal("kube-test-server-1"))
+		Eventually(responseFromInstance(multizone.KubeZone2), "30s", "1s").
+			Should(Equal("uni-test-server"))
 
 		// when
 		policy := `
@@ -109,7 +112,7 @@ spec:
   - targetRef:
       kind: MeshMultiZoneService
       labels:
-        kuma.io/display-name: test-server 
+        kuma.io/display-name: test-server
     default:
       localityAwareness:
         crossZone:
@@ -131,4 +134,35 @@ spec:
 		Eventually(responseFromInstance(multizone.KubeZone2), "30s", "1s").
 			MustPassRepeatedly(5).Should(Equal("kube-test-server-1"))
 	})
+
+	It("should be locality aware unless disabled", func() {
+		// given traffic only to the local zone
+		Eventually(responseFromInstance(multizone.KubeZone1), "30s", "1s").
+			MustPassRepeatedly(5).Should(Equal("kube-test-server-1"))
+
+		// when
+		policy := `
+type: MeshLoadBalancingStrategy
+name: mlb-mzms
+mesh: mlb-mzms
+spec:
+  targetRef:
+    kind: Mesh
+  to:
+  - targetRef:
+      kind: MeshMultiZoneService
+      labels:
+        kuma.io/display-name: test-server
+    default:
+      localityAwareness:
+        disabled: true
+`
+		err := multizone.Global.Install(YamlUniversal(policy))
+
+		// then
+		Expect(err).ToNot(HaveOccurred())
+
+		Eventually(responseFromInstance(multizone.KubeZone1), "30s", "1s").
+			Should(Equal("uni-test-server"))
+	})
 }