fix: verify the weight of the alb at the end of the rollout (#3627)

* fix: verify the weight of the alb at the end of the rollout when we auto set max weight Signed-off-by: Zach Aller <[email protected]> * add unit test Signed-off-by: Zach Aller <[email protected]> * refactor add unit test Signed-off-by: Zach Aller <[email protected]> * Trigger Build Signed-off-by: Zach Aller <[email protected]> * Trigger Build Signed-off-by: Zach Aller <[email protected]> * block reconcile Signed-off-by: Zach Aller <[email protected]> * add tests Signed-off-by: Zach Aller <[email protected]> * rename test and cleanup un-need logic check Signed-off-by: Zach Aller <[email protected]> --------- Signed-off-by: Zach Aller <[email protected]>
argoproj · Jun 13, 2024 · d28e75f · d28e75f
1 parent 02ab69e
commit d28e75f
Show file tree

Hide file tree

Showing 6 changed files with 112 additions and 8 deletions.
diff --git a/rollout/trafficrouting.go b/rollout/trafficrouting.go
@@ -286,6 +286,12 @@ func (c *rolloutContext) reconcileTrafficRouting() error {
 			} else {
 				c.log.Infof("Desired weight (stepIdx: %s) %d not yet verified", indexString, desiredWeight)
 				c.enqueueRolloutAfter(c.rollout, defaults.GetRolloutVerifyRetryInterval())
+				// At the end of the rollout we need to verify the weight is correct, and return an error if not because we don't want the rest of the
+				// reconcile process to continue. We don't need to do this if we are in the middle of the rollout because the rest of the reconcile
+				// process won't scale down the old replicasets yet due to being in the middle of some steps.
+				if desiredWeight == weightutil.MaxTrafficWeight(c.rollout) && len(c.rollout.Spec.Strategy.Canary.Steps) >= int(*c.rollout.Status.CurrentStepIndex) {
+					return fmt.Errorf("end of rollout, desired weight %d not yet verified", desiredWeight)
+				}
 			}
 		}
 	}

diff --git a/rollout/trafficrouting/alb/alb.go b/rollout/trafficrouting/alb/alb.go
@@ -202,7 +202,7 @@ func (r *Reconciler) VerifyWeight(desiredWeight int32, additionalDestinations ..
 		return nil, nil
 	}
 
-	if !rolloututil.ShouldVerifyWeight(r.cfg.Rollout) {
+	if !rolloututil.ShouldVerifyWeight(r.cfg.Rollout, desiredWeight) {
 		// If we should not verify weight but the ALB status has not been set yet due to a Rollout resource just being
 		// installed in the cluster we want to actually run the rest of the function, so we do not return if
 		// r.cfg.Rollout.Status.ALB is nil. However, if we should not verify, and we have already updated the status once

diff --git a/rollout/trafficrouting/alb/alb_test.go b/rollout/trafficrouting/alb/alb_test.go
@@ -911,6 +911,7 @@ func TestVerifyWeight(t *testing.T) {
 	{
 		var status v1alpha1.RolloutStatus
 		r, fakeClient := newFakeReconciler(&status)
+
 		fakeClient.loadBalancers = []*elbv2types.LoadBalancer{
 			{
 				LoadBalancerName: pointer.StringPtr("lb-abc123-name"),
@@ -948,6 +949,48 @@ func TestVerifyWeight(t *testing.T) {
 		assert.Equal(t, *status.ALB, *fakeClient.getAlbStatus("ingress"))
 	}
 
+	// LoadBalancer found, at max weight, end of rollout
+	{
+		var status v1alpha1.RolloutStatus
+		status.CurrentStepIndex = pointer.Int32Ptr(2)
+		r, fakeClient := newFakeReconciler(&status)
+		fakeClient.loadBalancers = []*elbv2types.LoadBalancer{
+			{
+				LoadBalancerName: pointer.StringPtr("lb-abc123-name"),
+				LoadBalancerArn:  pointer.StringPtr("arn:aws:elasticloadbalancing:us-east-2:123456789012:loadbalancer/app/lb-abc123-name/1234567890123456"),
+				DNSName:          pointer.StringPtr("verify-weight-test-abc-123.us-west-2.elb.amazonaws.com"),
+			},
+		}
+		fakeClient.targetGroups = []aws.TargetGroupMeta{
+			{
+				TargetGroup: elbv2types.TargetGroup{
+					TargetGroupName: pointer.StringPtr("canary-tg-abc123-name"),
+					TargetGroupArn:  pointer.StringPtr("arn:aws:elasticloadbalancing:us-east-2:123456789012:targetgroup/canary-tg-abc123-name/1234567890123456"),
+				},
+				Weight: pointer.Int32Ptr(100),
+				Tags: map[string]string{
+					aws.AWSLoadBalancerV2TagKeyResourceID: "default/ingress-canary-svc:443",
+				},
+			},
+			{
+				TargetGroup: elbv2types.TargetGroup{
+					TargetGroupName: pointer.StringPtr("stable-tg-abc123-name"),
+					TargetGroupArn:  pointer.StringPtr("arn:aws:elasticloadbalancing:us-east-2:123456789012:targetgroup/stable-tg-abc123-name/1234567890123456"),
+				},
+				Weight: pointer.Int32Ptr(0),
+				Tags: map[string]string{
+					aws.AWSLoadBalancerV2TagKeyResourceID: "default/ingress-stable-svc:443",
+				},
+			},
+		}
+
+		weightVerified, err := r.VerifyWeight(100)
+		assert.NoError(t, err)
+		assert.True(t, *weightVerified)
+		assert.Equal(t, status.ALBs[0], *status.ALB)
+		assert.Equal(t, *status.ALB, *fakeClient.getAlbStatus("ingress"))
+	}
+
 	// LoadBalancer found, but ARNs are unparsable
 	{
 		var status v1alpha1.RolloutStatus

diff --git a/rollout/trafficrouting_test.go b/rollout/trafficrouting_test.go
@@ -130,6 +130,53 @@ func TestReconcileTrafficRoutingVerifyWeightFalse(t *testing.T) {
 	assert.True(t, enqueued)
 }
 
+func TestReconcileTrafficRoutingVerifyWeightEndOfRollout(t *testing.T) {
+	f := newFixture(t)
+	defer f.Close()
+
+	steps := []v1alpha1.CanaryStep{
+		{
+			SetWeight: pointer.Int32Ptr(10),
+		},
+		{
+			Pause: &v1alpha1.RolloutPause{},
+		},
+	}
+	r1 := newCanaryRollout("foo", 10, nil, steps, pointer.Int32Ptr(2), intstr.FromInt(1), intstr.FromInt(0))
+	r2 := bumpVersion(r1)
+	r2.Spec.Strategy.Canary.TrafficRouting = &v1alpha1.RolloutTrafficRouting{}
+	r2.Spec.Strategy.Canary.CanaryService = "canary"
+	r2.Spec.Strategy.Canary.StableService = "stable"
+
+	rs1 := newReplicaSetWithStatus(r1, 10, 10)
+	rs2 := newReplicaSetWithStatus(r2, 10, 10)
+
+	rs1PodHash := rs1.Labels[v1alpha1.DefaultRolloutUniqueLabelKey]
+	rs2PodHash := rs2.Labels[v1alpha1.DefaultRolloutUniqueLabelKey]
+	canarySelector := map[string]string{v1alpha1.DefaultRolloutUniqueLabelKey: rs2PodHash}
+	stableSelector := map[string]string{v1alpha1.DefaultRolloutUniqueLabelKey: rs1PodHash}
+	canarySvc := newService("canary", 80, canarySelector, r2)
+	stableSvc := newService("stable", 80, stableSelector, r2)
+
+	f.kubeobjects = append(f.kubeobjects, rs1, rs2, canarySvc, stableSvc)
+	f.replicaSetLister = append(f.replicaSetLister, rs1, rs2)
+
+	r2 = updateCanaryRolloutStatus(r2, rs1PodHash, 10, 0, 10, false)
+	f.rolloutLister = append(f.rolloutLister, r2)
+	f.objects = append(f.objects, r2)
+
+	f.fakeTrafficRouting = newUnmockedFakeTrafficRoutingReconciler()
+	f.fakeTrafficRouting.On("UpdateHash", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	f.fakeTrafficRouting.On("SetWeight", mock.Anything, mock.Anything).Return(func(desiredWeight int32, additionalDestinations ...v1alpha1.WeightDestination) error {
+		// make sure SetWeight was called with correct value
+		assert.Equal(t, int32(100), desiredWeight)
+		return nil
+	})
+	f.fakeTrafficRouting.On("SetHeaderRoute", mock.Anything, mock.Anything).Return(nil)
+	f.fakeTrafficRouting.On("VerifyWeight", mock.Anything).Return(pointer.BoolPtr(false), nil)
+	f.runExpectError(getKey(r2, t), true)
+}
+
 func TestRolloutUseDesiredWeight(t *testing.T) {
 	f := newFixture(t)
 	defer f.Close()

diff --git a/utils/rollout/rolloututil.go b/utils/rollout/rolloututil.go
@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"strconv"
 
+	"github.com/argoproj/argo-rollouts/utils/weightutil"
+
 	replicasetutil "github.com/argoproj/argo-rollouts/utils/replicaset"
 
 	"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
@@ -184,13 +186,13 @@ func CanaryStepString(c v1alpha1.CanaryStep) string {
 
 // ShouldVerifyWeight We use this to test if we should verify weights because weight verification could involve
 // API calls to the cloud provider which could incur rate limiting
-func ShouldVerifyWeight(ro *v1alpha1.Rollout) bool {
+func ShouldVerifyWeight(ro *v1alpha1.Rollout, desiredWeight int32) bool {
 	currentStep, _ := replicasetutil.GetCurrentCanaryStep(ro)
 	// If we are in the middle of an update at a setWeight step, also perform weight verification.
 	// Note that we don't do this every reconciliation because weight verification typically involves
 	// API calls to the cloud provider which could incur rate limitingq
-	shouldVerifyWeight := ro.Status.StableRS != "" &&
-		!IsFullyPromoted(ro) &&
-		currentStep != nil && currentStep.SetWeight != nil
+	shouldVerifyWeight := (ro.Status.StableRS != "" && !IsFullyPromoted(ro) && currentStep != nil && currentStep.SetWeight != nil) ||
+		(ro.Status.StableRS != "" && !IsFullyPromoted(ro) && currentStep == nil && desiredWeight == weightutil.MaxTrafficWeight(ro)) // We are at end of rollout
+
 	return shouldVerifyWeight
 }
diff --git a/utils/rollout/rolloututil_test.go b/utils/rollout/rolloututil_test.go
@@ -422,15 +422,21 @@ func TestShouldVerifyWeight(t *testing.T) {
 	ro.Spec.Strategy.Canary.Steps = []v1alpha1.CanaryStep{{
 		SetWeight: pointer.Int32Ptr(20),
 	}}
-	assert.Equal(t, true, ShouldVerifyWeight(ro))
+	assert.Equal(t, true, ShouldVerifyWeight(ro, 20))
 
 	ro.Status.StableRS = ""
-	assert.Equal(t, false, ShouldVerifyWeight(ro))
+	assert.Equal(t, false, ShouldVerifyWeight(ro, 20))
 
 	ro.Status.StableRS = "34feab23f"
 	ro.Status.CurrentStepIndex = nil
 	ro.Spec.Strategy.Canary.Steps = nil
-	assert.Equal(t, false, ShouldVerifyWeight(ro))
+	assert.Equal(t, false, ShouldVerifyWeight(ro, 20))
+
+	// Test when the weight is 100, because we are at end of rollout
+	ro.Status.StableRS = "34feab23f"
+	ro.Status.CurrentStepIndex = nil
+	ro.Spec.Strategy.Canary.Steps = nil
+	assert.Equal(t, true, ShouldVerifyWeight(ro, 100))
 }
 
 func Test_isGenerationObserved(t *testing.T) {