From 86761a1582110e0280d98cb5ad0d37f49820c113 Mon Sep 17 00:00:00 2001
From: "nastassia.dailidava" <nastassia.dailidava@allegrogroup.com>
Date: Fri, 5 Apr 2024 17:31:08 +0200
Subject: [PATCH] #624 Set flat priority only for services with traffic
 splitting (added condition)

---
 .../endpoints/EnvoyEndpointsFactory.kt        | 49 +++++++++++--------
 .../endpoints/EnvoyEndpointsFactoryTest.kt    | 36 ++++++++++++++
 2 files changed, 65 insertions(+), 20 deletions(-)

diff --git a/envoy-control-core/src/main/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/resource/endpoints/EnvoyEndpointsFactory.kt b/envoy-control-core/src/main/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/resource/endpoints/EnvoyEndpointsFactory.kt
index 2338696a9..c95ddcbe9 100644
--- a/envoy-control-core/src/main/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/resource/endpoints/EnvoyEndpointsFactory.kt
+++ b/envoy-control-core/src/main/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/resource/endpoints/EnvoyEndpointsFactory.kt
@@ -100,30 +100,39 @@ class EnvoyEndpointsFactory(
         llbEndpointsList: List<LocalityLbEndpoints>,
         weights: ZoneWeights
     ): List<LocalityLbEndpoints> {
-        val endpoints = llbEndpointsList
-            .map {
-                if (weights.weightByZone.containsKey(it.locality.zone)) {
-                    LocalityLbEndpoints.newBuilder(it)
-                        .setLoadBalancingWeight(UInt32Value.of(weights.weightByZone[it.locality.zone] ?: 0))
-                        .build()
-                } else it
-            }
-        return overrideTrafficSplittingZoneEndpointsPriority(endpoints)
-            ?.let { listOf(it) + endpoints } ?: endpoints
+        if (properties.loadBalancing.trafficSplitting.zonesAllowingTrafficSplitting.contains(currentZone) &&
+            acceptableEndpointsRatioBetweenZones(llbEndpointsList)
+        ) {
+            val endpoints = llbEndpointsList
+                .map {
+                    if (weights.weightByZone.containsKey(it.locality.zone)) {
+                        LocalityLbEndpoints.newBuilder(it)
+                            .setLoadBalancingWeight(UInt32Value.of(weights.weightByZone[it.locality.zone] ?: 0))
+                            .build()
+                    } else it
+                }
+            return overrideTrafficSplittingZoneEndpointsPriority(endpoints) + endpoints
+        }
+        return llbEndpointsList
+    }
+
+    private fun acceptableEndpointsRatioBetweenZones(llbEndpointsList: List<LocalityLbEndpoints>): Boolean {
+        val localEndpoints = llbEndpointsList.filter { it.locality.zone == currentZone }
+        val trafficSplittingEndpoints = llbEndpointsList
+            .filter { it.locality.zone == properties.loadBalancing.trafficSplitting.zoneName }
+        return localEndpoints.isNotEmpty() && localEndpoints.size >= trafficSplittingEndpoints.size
     }
 
     private fun overrideTrafficSplittingZoneEndpointsPriority(
         endpoints: List<LocalityLbEndpoints>
-    ): LocalityLbEndpoints? {
-        return if (properties.loadBalancing.trafficSplitting.zonesAllowingTrafficSplitting.contains(currentZone)) {
-            endpoints
-                .find { properties.loadBalancing.trafficSplitting.zoneName == it.locality.zone }
-                ?.let {
-                    LocalityLbEndpoints.newBuilder(it)
-                        .setPriority(HIGHEST_PRIORITY)
-                        .build()
-                }
-        } else null
+    ): List<LocalityLbEndpoints> {
+        return endpoints
+            .filter { properties.loadBalancing.trafficSplitting.zoneName == it.locality.zone }
+            .map {
+                LocalityLbEndpoints.newBuilder(it)
+                    .setPriority(HIGHEST_PRIORITY)
+                    .build()
+            }
     }
 
     private fun filterEndpoints(loadAssignment: ClusterLoadAssignment, tag: String): ClusterLoadAssignment? {
diff --git a/envoy-control-core/src/test/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/resource/endpoints/EnvoyEndpointsFactoryTest.kt b/envoy-control-core/src/test/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/resource/endpoints/EnvoyEndpointsFactoryTest.kt
index de151d627..9e953b531 100644
--- a/envoy-control-core/src/test/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/resource/endpoints/EnvoyEndpointsFactoryTest.kt
+++ b/envoy-control-core/src/test/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/resource/endpoints/EnvoyEndpointsFactoryTest.kt
@@ -406,6 +406,42 @@ internal class EnvoyEndpointsFactoryTest {
             .anySatisfy { it.hasZoneWithPriority("DC3", 2) }
     }
 
+    @Test
+    fun `should not duplicate endpoints if endpoints ratio is lower than acceptable threshold`() {
+        val envoyEndpointsFactory = EnvoyEndpointsFactory(
+            snapshotPropertiesWithWeights,
+            currentZone = "DC1"
+        )
+        val loadAssignments = envoyEndpointsFactory.createLoadAssignment(
+            setOf(serviceName),
+            MultiClusterState(
+                listOf(
+                    clusterState(Locality.LOCAL, "DC1"),
+                    ClusterState(
+                        ServicesState(
+                            concurrentMapOf(
+                                serviceName to ServiceInstances(serviceName, setOf())
+                            )
+                        ), Locality.REMOTE, "DC2"
+                    ),
+                    clusterState(Locality.REMOTE, "DC3")
+                )
+            )
+        )
+        var resultLoadAssignment = envoyEndpointsFactory.assignLocalityWeights(
+            WeightRouteSpecification(
+                serviceName, listOf(), DependencySettings(), ZoneWeights().apply { weightByZone = defaultZoneWeights }
+            ),
+            loadAssignments[0]
+        )
+
+        assertThat(resultLoadAssignment.endpointsList).hasSize(dcPriorityProperties.size + 1)
+        assertThat(resultLoadAssignment.endpointsList)
+            .anySatisfy { it.hasZoneWithPriority("DC2", 1) }
+            .anySatisfy { it.hasZoneWithPriority("DC1", 0) }
+            .anySatisfy { it.hasZoneWithPriority("DC3", 2) }
+    }
+
     private fun List<ClusterLoadAssignment>.assertHasLoadAssignment(map: Map<String, Int>) {
         assertThat(this)
             .isNotEmpty()