feat: TrafficRouting SMI with Experiment Step in Canary (#1351)

* feat: Allow Experiment step in TrafficRouting for SMI

Signed-off-by: khhirani <[email protected]>

hack/update-mocks.sh

@@ -18,6 +18,6 @@ mockery \
     --output "${PROJECT_ROOT}"/utils/aws/mocks
 
 mockery \
-    --dir "${PROJECT_ROOT}"/rollout \
+    --dir "${PROJECT_ROOT}"/rollout/trafficrouting \
     --name TrafficRoutingReconciler \
     --output "${PROJECT_ROOT}"/rollout/mocks

manifests/crds/rollout-crd.yaml

@@ -459,6 +459,9 @@ spec:
                                         type: object
                                       specRef:
                                         type: string
+                                      weight:
+                                        format: int32
+                                        type: integer
                                     required:
                                     - name
                                     - specRef

manifests/install.yaml

@@ -10157,6 +10157,9 @@ spec:
                                         type: object
                                       specRef:
                                         type: string
+                                      weight:
+                                        format: int32
+                                        type: integer
                                     required:
                                     - name
                                     - specRef

manifests/namespace-install.yaml

@@ -10157,6 +10157,9 @@ spec:
                                         type: object
                                       specRef:
                                         type: string
+                                      weight:
+                                        format: int32
+                                        type: integer
                                     required:
                                     - name
                                     - specRef

pkg/apiclient/rollout/rollout.swagger.json

@@ -1128,6 +1128,11 @@
         "selector": {
           "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.LabelSelector",
           "title": "Selector overrides the selector to be used for the template's ReplicaSet. If omitted, will\nuse the same selector as the Rollout\n+optional"
+        },
+        "weight": {
+          "type": "integer",
+          "format": "int32",
+          "title": "Weight sets the percentage of traffic the template's replicas should receive"
         }
       },
       "title": "RolloutExperimentTemplate defines the template used to create experiments for the Rollout's experiment canary step"

pkg/apis/rollouts/v1alpha1/types.go

@@ -449,6 +449,8 @@ type RolloutExperimentTemplate struct {
 	// use the same selector as the Rollout
 	// +optional
 	Selector *metav1.LabelSelector `json:"selector,omitempty" protobuf:"bytes,5,opt,name=selector"`
+	// Weight sets the percentage of traffic the template's replicas should receive
+	Weight *int32 `json:"weight,omitempty" protobuf:"varint,6,opt,name=weight"`
 }
 
 // PodTemplateMetadata extra labels to add to the template

rollout/controller.go

@@ -40,6 +40,7 @@ import (
 	clientset "github.com/argoproj/argo-rollouts/pkg/client/clientset/versioned"
 	informers "github.com/argoproj/argo-rollouts/pkg/client/informers/externalversions/rollouts/v1alpha1"
 	listers "github.com/argoproj/argo-rollouts/pkg/client/listers/rollouts/v1alpha1"
+	"github.com/argoproj/argo-rollouts/rollout/trafficrouting"
 	"github.com/argoproj/argo-rollouts/rollout/trafficrouting/ambassador"
 	"github.com/argoproj/argo-rollouts/rollout/trafficrouting/istio"
 	analysisutil "github.com/argoproj/argo-rollouts/utils/analysis"
@@ -138,9 +139,9 @@ type reconcilerBase struct {
 	podRestarter RolloutPodRestarter
 
 	// used for unit testing
-	enqueueRollout              func(obj interface{})                                         //nolint:structcheck
-	enqueueRolloutAfter         func(obj interface{}, duration time.Duration)                 //nolint:structcheck
-	newTrafficRoutingReconciler func(roCtx *rolloutContext) (TrafficRoutingReconciler, error) //nolint:structcheck
+	enqueueRollout              func(obj interface{})                                                        //nolint:structcheck
+	enqueueRolloutAfter         func(obj interface{}, duration time.Duration)                                //nolint:structcheck
+	newTrafficRoutingReconciler func(roCtx *rolloutContext) (trafficrouting.TrafficRoutingReconciler, error) //nolint:structcheck
 
 	// recorder is an event recorder for recording Event resources to the Kubernetes API.
 	recorder     record.EventRecorder

rollout/controller_test.go

@@ -46,6 +46,7 @@ import (
 	"github.com/argoproj/argo-rollouts/pkg/client/clientset/versioned/fake"
 	informers "github.com/argoproj/argo-rollouts/pkg/client/informers/externalversions"
 	"github.com/argoproj/argo-rollouts/rollout/mocks"
+	"github.com/argoproj/argo-rollouts/rollout/trafficrouting"
 	"github.com/argoproj/argo-rollouts/utils/annotations"
 	"github.com/argoproj/argo-rollouts/utils/conditions"
 	"github.com/argoproj/argo-rollouts/utils/defaults"
@@ -539,7 +540,7 @@ func (f *fixture) newController(resync resyncFunc) (*Controller, informers.Share
 		c.enqueueRollout(obj)
 	}
 
-	c.newTrafficRoutingReconciler = func(roCtx *rolloutContext) (TrafficRoutingReconciler, error) {
+	c.newTrafficRoutingReconciler = func(roCtx *rolloutContext) (trafficrouting.TrafficRoutingReconciler, error) {
 		return f.fakeTrafficRouting, nil
 	}
 

rollout/experiment.go

@@ -63,6 +63,9 @@ func GetExperimentFromTemplate(r *v1alpha1.Rollout, stableRS, newRS *appsv1.Repl
 			Name:     templateStep.Name,
 			Replicas: templateStep.Replicas,
 		}
+		if templateStep.Weight != nil {
+			template.Service = &v1alpha1.TemplateService{}
+		}
 		templateRS := &appsv1.ReplicaSet{}
 		switch templateStep.SpecRef {
 		case v1alpha1.CanarySpecRef:

rollout/experiment_test.go

@@ -750,3 +750,46 @@ func TestRolloutCreateExperimentWithInstanceID(t *testing.T) {
 	assert.Equal(t, createdEx.Name, ex.Name)
 	assert.Equal(t, "instance-id-test", createdEx.Labels[v1alpha1.LabelKeyControllerInstanceID])
 }
+
+// TestRolloutCreateExperimentWithService verifies the controller sets CreateService for Experiment Template as expected.
+// CreateService is true when Weight is set in RolloutExperimentStep for template, otherwise false.
+func TestRolloutCreateExperimentWithService(t *testing.T) {
+	steps := []v1alpha1.CanaryStep{{
+		Experiment: &v1alpha1.RolloutExperimentStep{
+			Templates: []v1alpha1.RolloutExperimentTemplate{
+				// Service should be created for "stable-template"
+				{
+					Name:     "stable-template",
+					SpecRef:  v1alpha1.StableSpecRef,
+					Replicas: pointer.Int32Ptr(1),
+					Weight:   pointer.Int32Ptr(5),
+				},
+				// Service should NOT be created for "canary-template"
+				{
+					Name:     "canary-template",
+					SpecRef:  v1alpha1.CanarySpecRef,
+					Replicas: pointer.Int32Ptr(1),
+				},
+			},
+		},
+	}}
+
+	r1 := newCanaryRollout("foo", 1, nil, steps, pointer.Int32Ptr(0), intstr.FromInt(0), intstr.FromInt(1))
+	r2 := bumpVersion(r1)
+
+	rs1 := newReplicaSetWithStatus(r1, 1, 1)
+	rs2 := newReplicaSetWithStatus(r2, 1, 1)
+	rs1PodHash := rs1.Labels[v1alpha1.DefaultRolloutUniqueLabelKey]
+
+	r2.Status.CurrentStepIndex = pointer.Int32Ptr(0)
+	r2.Status.StableRS = rs1PodHash
+
+	ex, err := GetExperimentFromTemplate(r2, rs1, rs2)
+	assert.Nil(t, err)
+
+	assert.Equal(t, "stable-template", ex.Spec.Templates[0].Name)
+	assert.NotNil(t, ex.Spec.Templates[0].Service)
+
+	assert.Equal(t, "canary-template", ex.Spec.Templates[1].Name)
+	assert.Nil(t, ex.Spec.Templates[1].Service)
+}

rollout/trafficrouting.go

@@ -10,24 +10,13 @@ import (
 	"github.com/argoproj/argo-rollouts/rollout/trafficrouting/nginx"
 	"github.com/argoproj/argo-rollouts/rollout/trafficrouting/smi"
 
+	"github.com/argoproj/argo-rollouts/rollout/trafficrouting"
 	"github.com/argoproj/argo-rollouts/utils/record"
 	replicasetutil "github.com/argoproj/argo-rollouts/utils/replicaset"
 )
 
-// TrafficRoutingReconciler common function across all TrafficRouting implementation
-type TrafficRoutingReconciler interface {
-	// UpdateHash informs a traffic routing reconciler about new canary/stable pod hashes
-	UpdateHash(canaryHash, stableHash string) error
-	// SetWeight sets the canary weight to the desired weight
-	SetWeight(desiredWeight int32) error
-	// VerifyWeight returns true if the canary is at the desired weight
-	VerifyWeight(desiredWeight int32) (bool, error)
-	// Type returns the type of the traffic routing reconciler
-	Type() string
-}
-
 // NewTrafficRoutingReconciler identifies return the TrafficRouting Plugin that the rollout wants to modify
-func (c *Controller) NewTrafficRoutingReconciler(roCtx *rolloutContext) (TrafficRoutingReconciler, error) {
+func (c *Controller) NewTrafficRoutingReconciler(roCtx *rolloutContext) (trafficrouting.TrafficRoutingReconciler, error) {
 	rollout := roCtx.rollout
 	if rollout.Spec.Strategy.Canary.TrafficRouting == nil {
 		return nil, nil
@@ -96,6 +85,7 @@ func (c *rolloutContext) reconcileTrafficRouting() error {
 
 	currentStep, index := replicasetutil.GetCurrentCanaryStep(c.rollout)
 	desiredWeight := int32(0)
+	weightDestinations := make([]trafficrouting.WeightDestination, 0)
 	if c.rollout.Status.StableRS == c.rollout.Status.CurrentPodHash {
 		// when we are fully promoted. desired canary weight should be 0
 	} else if c.pauseContext.IsAborted() {
@@ -122,9 +112,31 @@ func (c *rolloutContext) reconcileTrafficRouting() error {
 			// last setWeight step, which is set by GetCurrentSetWeight.
 			desiredWeight = replicasetutil.GetCurrentSetWeight(c.rollout)
 		}
+
+		// Checks for experiment step
+		// If current experiment exists, then create WeightDestinations for each experiment template
+		exStep := replicasetutil.GetCurrentExperimentStep(c.rollout)
+		if exStep != nil && c.currentEx != nil && c.currentEx.Status.Phase == v1alpha1.AnalysisPhaseRunning {
+			getTemplateWeight := func(name string) *int32 {
+				for _, tmpl := range exStep.Templates {
+					if tmpl.Name == name {
+						return tmpl.Weight
+					}
+				}
+				return nil
+			}
+			for _, templateStatus := range c.currentEx.Status.TemplateStatuses {
+				templateWeight := getTemplateWeight(templateStatus.Name)
+				weightDestinations = append(weightDestinations, trafficrouting.WeightDestination{
+					ServiceName:     templateStatus.ServiceName,
+					PodTemplateHash: templateStatus.PodTemplateHash,
+					Weight:          *templateWeight,
+				})
+			}
+		}
 	}
 
-	err = reconciler.SetWeight(desiredWeight)
+	err = reconciler.SetWeight(desiredWeight, weightDestinations...)
 	if err != nil {
 		c.recorder.Warnf(c.rollout, record.EventOptions{EventReason: "TrafficRoutingError"}, err.Error())
 		return err

rollout/trafficrouting/alb/alb.go

@@ -15,6 +15,7 @@ import (
 	"k8s.io/utils/pointer"
 
 	"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
+	"github.com/argoproj/argo-rollouts/rollout/trafficrouting"
 	"github.com/argoproj/argo-rollouts/utils/aws"
 	"github.com/argoproj/argo-rollouts/utils/diff"
 	ingressutil "github.com/argoproj/argo-rollouts/utils/ingress"
@@ -74,7 +75,7 @@ func (r *Reconciler) Type() string {
 }
 
 // SetWeight modifies ALB Ingress resources to reach desired state
-func (r *Reconciler) SetWeight(desiredWeight int32) error {
+func (r *Reconciler) SetWeight(desiredWeight int32, additionalDestinations ...trafficrouting.WeightDestination) error {
 	ctx := context.TODO()
 	rollout := r.cfg.Rollout
 	ingressName := rollout.Spec.Strategy.Canary.TrafficRouting.ALB.Ingress

rollout/trafficrouting/ambassador/ambassador.go

@@ -18,6 +18,7 @@ import (
 	"k8s.io/client-go/dynamic"
 
 	"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
+	"github.com/argoproj/argo-rollouts/rollout/trafficrouting"
 	"github.com/argoproj/argo-rollouts/utils/defaults"
 	logutil "github.com/argoproj/argo-rollouts/utils/log"
 	"github.com/argoproj/argo-rollouts/utils/record"
@@ -88,7 +89,7 @@ func NewReconciler(r *v1alpha1.Rollout, c ClientInterface, rec record.EventRecor
 // in the ambassador configuration in the traffic routing section of the rollout. If
 // the canary ambassador mapping is already present, it will be updated to the given
 // desiredWeight.
-func (r *Reconciler) SetWeight(desiredWeight int32) error {
+func (r *Reconciler) SetWeight(desiredWeight int32, additionalDestinations ...trafficrouting.WeightDestination) error {
 	r.sendNormalEvent(CanaryMappingWeightUpdate, fmt.Sprintf("Set canary mapping weight to %d", desiredWeight))
 	ctx := context.TODO()
 	baseMappingNameList := r.Rollout.Spec.Strategy.Canary.TrafficRouting.Ambassador.Mappings

rollout/trafficrouting/ambassador/ambassador_test.go

@@ -7,16 +7,17 @@ import (
 	"sync"
 	"testing"
 
-	"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
-	"github.com/argoproj/argo-rollouts/rollout/trafficrouting/ambassador"
-	"github.com/argoproj/argo-rollouts/utils/record"
 	"github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/assert"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/runtime/serializer/yaml"
+
+	"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
+	"github.com/argoproj/argo-rollouts/rollout/trafficrouting/ambassador"
+	"github.com/argoproj/argo-rollouts/utils/record"
 )
 
 const (

rollout/trafficrouting/istio/controller.go

@@ -5,8 +5,6 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/argoproj/argo-rollouts/utils/queue"
-
 	log "github.com/sirupsen/logrus"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/meta"
@@ -28,6 +26,7 @@ import (
 	"github.com/argoproj/argo-rollouts/utils/defaults"
 	istioutil "github.com/argoproj/argo-rollouts/utils/istio"
 	logutil "github.com/argoproj/argo-rollouts/utils/log"
+	"github.com/argoproj/argo-rollouts/utils/queue"
 	unstructuredutil "github.com/argoproj/argo-rollouts/utils/unstructured"
 )
 

rollout/trafficrouting/istio/istio.go

@@ -16,6 +16,7 @@ import (
 	"k8s.io/client-go/dynamic/dynamiclister"
 
 	"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
+	"github.com/argoproj/argo-rollouts/rollout/trafficrouting"
 	evalUtils "github.com/argoproj/argo-rollouts/utils/evaluate"
 	istioutil "github.com/argoproj/argo-rollouts/utils/istio"
 	logutil "github.com/argoproj/argo-rollouts/utils/log"
@@ -522,7 +523,7 @@ func (r *Reconciler) Type() string {
 }
 
 // SetWeight modifies Istio resources to reach desired state
-func (r *Reconciler) SetWeight(desiredWeight int32) error {
+func (r *Reconciler) SetWeight(desiredWeight int32, additionalDestinations ...trafficrouting.WeightDestination) error {
 	ctx := context.TODO()
 	var vsvc *unstructured.Unstructured
 	var err error