Added response header for traffic splitting feature

allegro · Dec 7, 2023 · dae5a0a · dae5a0a
1 parent 1abb570
commit dae5a0a
Show file tree

Hide file tree

Showing 4 changed files with 100 additions and 2 deletions.
diff --git a/...e/src/main/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/SnapshotProperties.kt b/...e/src/main/kotlin/pl/allegro/tech/servicemesh/envoycontrol/snapshot/SnapshotProperties.kt
@@ -157,6 +157,7 @@ class CanaryProperties {
 
 class TrafficSplittingProperties {
     var zoneName = ""
+    var headerName = ""
     var serviceByWeightsProperties: Map<String, ZoneWeights> = mapOf()
     var secondaryClusterSuffix = "secondary"
     var aggregateClusterSuffix = "aggregate"

diff --git a/...llegro/tech/servicemesh/envoycontrol/snapshot/resource/routes/EnvoyEgressRoutesFactory.kt b/...llegro/tech/servicemesh/envoycontrol/snapshot/resource/routes/EnvoyEgressRoutesFactory.kt
@@ -359,7 +359,8 @@ class EnvoyEgressRoutesFactory(
                         .withClusterWeight(routeSpec.clusterName, routeSpec.clusterWeights.main)
                         .withClusterWeight(
                             getAggregateClusterName(routeSpec.clusterName, properties),
-                            routeSpec.clusterWeights.secondary
+                            routeSpec.clusterWeights.secondary,
+                            routeSpec.clusterName
                         )
                 )
             }
@@ -369,15 +370,42 @@ class EnvoyEgressRoutesFactory(
         }
     }
 
-    private fun WeightedCluster.Builder.withClusterWeight(clusterName: String, weight: Int): WeightedCluster.Builder {
+    private fun WeightedCluster.Builder.withClusterWeight(
+        clusterName: String,
+        weight: Int,
+        headerValue: String? = null
+    ): WeightedCluster.Builder {
         this.addClusters(
             WeightedCluster.ClusterWeight.newBuilder()
                 .setName(clusterName)
                 .setWeight(UInt32Value.of(weight))
+                .withHeader(properties.loadBalancing.trafficSplitting.headerName, headerValue)
                 .build()
         )
         return this
     }
+
+    private fun WeightedCluster.ClusterWeight.Builder.withHeader(
+        key: String?,
+        value: String?
+    ): WeightedCluster.ClusterWeight.Builder {
+        key?.takeIf { it.isNotBlank() }
+            ?.let {
+                value?.let { this.addResponseHeadersToAdd(buildHeader(key, value)) }
+            }
+        return this
+    }
+
+    private fun buildHeader(key: String, value: String): HeaderValueOption.Builder {
+        return HeaderValueOption.newBuilder()
+            .setHeader(
+                HeaderValue.newBuilder()
+                    .setKey(key)
+                    .setValue(value)
+            )
+            .setAppendAction(HeaderValueOption.HeaderAppendAction.OVERWRITE_IF_EXISTS_OR_ADD)
+            .setKeepEmptyValue(false)
+    }
 }
 
 class RequestPolicyMapper private constructor() {

diff --git a/...-core/src/test/kotlin/pl/allegro/tech/servicemesh/envoycontrol/groups/RoutesAssertions.kt b/...-core/src/test/kotlin/pl/allegro/tech/servicemesh/envoycontrol/groups/RoutesAssertions.kt
@@ -12,6 +12,7 @@ import io.envoyproxy.envoy.config.route.v3.RouteAction
 import io.envoyproxy.envoy.config.route.v3.RouteConfiguration
 import io.envoyproxy.envoy.config.route.v3.VirtualCluster
 import io.envoyproxy.envoy.config.route.v3.VirtualHost
+import io.envoyproxy.envoy.config.route.v3.WeightedCluster
 import io.envoyproxy.envoy.type.matcher.v3.RegexMatcher
 import org.assertj.core.api.Assertions.assertThat
 import pl.allegro.tech.servicemesh.envoycontrol.snapshot.LocalRetryPolicyProperties
@@ -27,6 +28,16 @@ fun RouteConfiguration.hasVirtualHostThat(name: String, condition: VirtualHost.(
     return this
 }
 
+fun RouteConfiguration.hasClusterThat(name: String, condition: WeightedCluster.ClusterWeight?.() -> Unit): RouteConfiguration {
+    condition(this.virtualHostsList
+        .flatMap { it.routesList }
+        .map { it.route }
+        .flatMap { route -> route.weightedClusters.clustersList }
+        .find { it.name == name }
+    )
+    return this
+}
+
 fun RouteConfiguration.hasRequestHeaderToAdd(key: String, value: String): RouteConfiguration {
     assertThat(this.requestHeadersToAddList).anySatisfy {
         assertThat(it.header.key).isEqualTo(key)

diff --git a/...ro/tech/servicemesh/envoycontrol/snapshot/resource/routes/EnvoyEgressRoutesFactoryTest.kt b/...ro/tech/servicemesh/envoycontrol/snapshot/resource/routes/EnvoyEgressRoutesFactoryTest.kt
@@ -2,10 +2,13 @@ package pl.allegro.tech.servicemesh.envoycontrol.snapshot.resource.routes
 
 import com.google.protobuf.util.Durations
 import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNotNull
+import org.junit.jupiter.api.Assertions.assertTrue
 import org.junit.jupiter.api.Test
 import pl.allegro.tech.servicemesh.envoycontrol.groups.DependencySettings
 import pl.allegro.tech.servicemesh.envoycontrol.groups.Outgoing
 import pl.allegro.tech.servicemesh.envoycontrol.groups.RetryPolicy
+import pl.allegro.tech.servicemesh.envoycontrol.groups.hasClusterThat
 import pl.allegro.tech.servicemesh.envoycontrol.groups.hasCustomIdleTimeout
 import pl.allegro.tech.servicemesh.envoycontrol.groups.hasCustomRequestTimeout
 import pl.allegro.tech.servicemesh.envoycontrol.groups.hasHostRewriteHeader
@@ -23,6 +26,8 @@ import pl.allegro.tech.servicemesh.envoycontrol.groups.matchingOnMethod
 import pl.allegro.tech.servicemesh.envoycontrol.groups.matchingOnPrefix
 import pl.allegro.tech.servicemesh.envoycontrol.snapshot.SnapshotProperties
 import pl.allegro.tech.servicemesh.envoycontrol.snapshot.StandardRouteSpecification
+import pl.allegro.tech.servicemesh.envoycontrol.snapshot.WeightRouteSpecification
+import pl.allegro.tech.servicemesh.envoycontrol.snapshot.ZoneWeights
 
 internal class EnvoyEgressRoutesFactoryTest {
 
@@ -42,6 +47,15 @@ internal class EnvoyEgressRoutesFactoryTest {
         )
     )
 
+    val weightedClusters = listOf(
+        WeightRouteSpecification(
+            clusterName = "srv1",
+            routeDomains = listOf("srv1"),
+            settings = DependencySettings(),
+            ZoneWeights()
+        )
+    )
+
     @Test
     fun `should add client identity header if incoming permissions are enabled`() {
         // given
@@ -278,4 +292,48 @@ internal class EnvoyEgressRoutesFactoryTest {
             defaultRoute.matchingOnPrefix("/")
         }
     }
+
+    @Test
+    fun `should add traffic splitting header for secondary weighted cluster`() {
+
+        val expectedHeaderKey = "test-header"
+        val routesFactory = EnvoyEgressRoutesFactory(SnapshotProperties().apply {
+            loadBalancing.trafficSplitting.headerName = expectedHeaderKey
+        })
+
+        val routeConfig = routesFactory.createEgressRouteConfig(
+            "client1",
+            weightedClusters,
+            false
+        )
+
+        routeConfig
+            .hasClusterThat("srv1-aggregate") {
+                assertNotNull(this)
+                assertNotNull(this!!.responseHeadersToAddList.find { it.header.key == expectedHeaderKey })
+            }.hasClusterThat("srv1") {
+                assertNotNull(this)
+                assertTrue(this!!.responseHeadersToAddList.isEmpty())
+            }
+    }
+
+
+    @Test
+    fun `should not add traffic splitting header if header key is not set`() {
+        val routesFactory = EnvoyEgressRoutesFactory(SnapshotProperties())
+        val routeConfig = routesFactory.createEgressRouteConfig(
+            "client1",
+            weightedClusters,
+            false
+        )
+
+        routeConfig
+            .hasClusterThat("srv1-aggregate") {
+                assertNotNull(this)
+                assertTrue(this!!.responseHeadersToAddList.isEmpty())
+            }.hasClusterThat("srv1") {
+                assertNotNull(this)
+                assertNotNull(this!!.responseHeadersToAddList.isEmpty())
+            }
+    }
 }