feat: support multiple backend refs (#53)

Signed-off-by: Tommy Chen <[email protected]>
argoproj-labs · Apr 29, 2024 · 47954c2 · 47954c2
1 parent b2477b9
commit 47954c2
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 10 deletions.
diff --git a/pkg/plugin/httproute.go b/pkg/plugin/httproute.go
@@ -32,21 +32,25 @@ func (r *RpcPlugin) setHTTPRouteWeight(rollout *v1alpha1.Rollout, desiredWeight
 	canaryServiceName := rollout.Spec.Strategy.Canary.CanaryService
 	stableServiceName := rollout.Spec.Strategy.Canary.StableService
 	routeRuleList := HTTPRouteRuleList(httpRoute.Spec.Rules)
-	canaryBackendRef, err := getBackendRef(canaryServiceName, routeRuleList)
+	canaryBackendRefs, err := getBackendRefs(canaryServiceName, routeRuleList)
 	if err != nil {
 		return pluginTypes.RpcError{
 			ErrorString: err.Error(),
 		}
 	}
-	canaryBackendRef.Weight = &desiredWeight
-	stableBackendRef, err := getBackendRef(stableServiceName, routeRuleList)
+	for _, ref := range canaryBackendRefs {
+		ref.Weight = &desiredWeight
+	}
+	stableBackendRefs, err := getBackendRefs(stableServiceName, routeRuleList)
 	if err != nil {
 		return pluginTypes.RpcError{
 			ErrorString: err.Error(),
 		}
 	}
 	restWeight := 100 - desiredWeight
-	stableBackendRef.Weight = &restWeight
+	for _, ref := range stableBackendRefs {
+		ref.Weight = &restWeight
+	}
 	updatedHTTPRoute, err := httpRouteClient.Update(ctx, httpRoute, metav1.UpdateOptions{})
 	if r.IsTest {
 		r.UpdatedHTTPRouteMock = updatedHTTPRoute

diff --git a/pkg/plugin/plugin.go b/pkg/plugin/plugin.go
@@ -201,18 +201,22 @@ func getRouteRule[T1 GatewayAPIBackendRef, T2 GatewayAPIRouteRule[T1], T3 Gatewa
 	return nil, routeRuleList.Error()
 }
 
-func getBackendRef[T1 GatewayAPIBackendRef, T2 GatewayAPIRouteRule[T1], T3 GatewayAPIRouteRuleList[T1, T2]](backendRefName string, routeRuleList T3) (T1, error) {
+func getBackendRefs[T1 GatewayAPIBackendRef, T2 GatewayAPIRouteRule[T1], T3 GatewayAPIRouteRuleList[T1, T2]](backendRefName string, routeRuleList T3) ([]T1, error) {
 	var backendRef T1
 	var routeRule T2
+	var matchedRefs []T1
 	for next, hasNext := routeRuleList.Iterator(); hasNext; {
 		routeRule, hasNext = next()
 		for next, hasNext := routeRule.Iterator(); hasNext; {
 			backendRef, hasNext = next()
 			if backendRefName == backendRef.GetName() {
-				return backendRef, nil
+				matchedRefs = append(matchedRefs, backendRef)
 			}
 		}
 	}
+	if len(matchedRefs) > 0 {
+		return matchedRefs, nil
+	}
 	return nil, routeRuleList.Error()
 }
 

diff --git a/pkg/plugin/tcproute.go b/pkg/plugin/tcproute.go
@@ -26,21 +26,25 @@ func (r *RpcPlugin) setTCPRouteWeight(rollout *v1alpha1.Rollout, desiredWeight i
 	canaryServiceName := rollout.Spec.Strategy.Canary.CanaryService
 	stableServiceName := rollout.Spec.Strategy.Canary.StableService
 	routeRuleList := TCPRouteRuleList(tcpRoute.Spec.Rules)
-	canaryBackendRef, err := getBackendRef(canaryServiceName, routeRuleList)
+	canaryBackendRefs, err := getBackendRefs(canaryServiceName, routeRuleList)
 	if err != nil {
 		return pluginTypes.RpcError{
 			ErrorString: err.Error(),
 		}
 	}
-	canaryBackendRef.Weight = &desiredWeight
-	stableBackendRef, err := getBackendRef(stableServiceName, routeRuleList)
+	for _, ref := range canaryBackendRefs {
+		ref.Weight = &desiredWeight
+	}
+	stableBackendRefs, err := getBackendRefs(stableServiceName, routeRuleList)
 	if err != nil {
 		return pluginTypes.RpcError{
 			ErrorString: err.Error(),
 		}
 	}
 	restWeight := 100 - desiredWeight
-	stableBackendRef.Weight = &restWeight
+	for _, ref := range stableBackendRefs {
+		ref.Weight = &restWeight
+	}
 	updatedTCPRoute, err := tcpRouteClient.Update(ctx, tcpRoute, metav1.UpdateOptions{})
 	if r.IsTest {
 		r.UpdatedTCPRouteMock = updatedTCPRoute