Automatically scale down Deployment after migrating to Rollout

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

manifests/crds/rollout-crd.yaml

@@ -3330,6 +3330,8 @@ spec:
                     type: string
                   name:
                     type: string
+                  scaleDown:
+                    type: string
                 type: object
             type: object
           status:

manifests/install.yaml

@@ -14920,6 +14920,8 @@ spec:
                     type: string
                   name:
                     type: string
+                  scaleDown:
+                    type: string
                 type: object
             type: object
           status:

pkg/apiclient/rollout/rollout.swagger.json

@@ -1695,6 +1695,10 @@
         "name": {
           "type": "string",
           "title": "Name of the referent"
+        },
+        "scaleDown": {
+          "type": "string",
+          "title": "Automatically scale down deployment"
         }
       },
       "title": "ObjectRef holds a references to the Kubernetes object"

pkg/apis/rollouts/v1alpha1/types.go

@@ -137,6 +137,8 @@ type ObjectRef struct {
 	Kind string `json:"kind,omitempty" protobuf:"bytes,2,opt,name=kind"`
 	// Name of the referent
 	Name string `json:"name,omitempty" protobuf:"bytes,3,opt,name=name"`
+	// Automatically scale down deployment
+	ScaleDown string `json:"scaleDown,omitempty" protobuf:"bytes,4,opt,name=scaleDown"`
 }
 
 const (
@@ -1081,3 +1083,9 @@ type RolloutList struct {
 type RollbackWindowSpec struct {
 	Revisions int32 `json:"revisions,omitempty" protobuf:"varint,1,opt,name=revisions"`
 }
+
+const (
+	ScaleDownNever         string = "never"
+	ScaleDownOnSuccess     string = "onsuccess"
+	ScaleDownProgressively string = "progressively"
+)

rollout/sync.go

@@ -378,6 +378,12 @@ func (c *rolloutContext) scaleReplicaSet(rs *appsv1.ReplicaSet, newScale int32,
 			revision, _ := replicasetutil.Revision(rs)
 			c.recorder.Eventf(rollout, record.EventOptions{EventReason: conditions.ScalingReplicaSetReason}, conditions.ScalingReplicaSetMessage, scalingOperation, rs.Name, revision, oldScale, newScale)
 			c.replicaSetInformer.GetIndexer().Update(rs)
+			if rollout.Spec.WorkloadRef != nil && rollout.Spec.WorkloadRef.ScaleDown == v1alpha1.ScaleDownProgressively {
+				err = c.scaleDeployment(false, &oldScale, &newScale)
+				if err != nil {
+					return scaled, rs, err
+				}
+			}
 		}
 	}
 	return scaled, rs, err
@@ -1003,6 +1009,34 @@ func (c *rolloutContext) promoteStable(newStatus *v1alpha1.RolloutStatus, reason
 		revision, _ := replicasetutil.Revision(c.rollout)
 		c.recorder.Eventf(c.rollout, record.EventOptions{EventReason: conditions.RolloutCompletedReason},
 			conditions.RolloutCompletedMessage, revision, newStatus.CurrentPodHash, reason)
+
+		if c.rollout.Spec.WorkloadRef != nil && c.rollout.Spec.WorkloadRef.ScaleDown == v1alpha1.ScaleDownOnSuccess {
+			err := c.scaleDeployment(true, nil, nil)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (c *rolloutContext) scaleDeployment(scaleToZero bool, oldScale *int32, newScale *int32) error {
+	deploymentName := c.rollout.Spec.WorkloadRef.Name
+	namespace := c.rollout.Namespace
+	deployment, err := c.kubeclientset.AppsV1().Deployments(namespace).Get(context.TODO(), deploymentName, metav1.GetOptions{})
+	if err != nil {
+		c.log.Warnf("Failed to fetch deployment %s: %s", deploymentName, err.Error())
+		return err
+	}
+	if scaleToZero {
+		*deployment.Spec.Replicas = 0
+	} else {
+		*deployment.Spec.Replicas += *oldScale - *newScale
+	}
+	_, err = c.kubeclientset.AppsV1().Deployments(namespace).Update(context.TODO(), deployment, metav1.UpdateOptions{})
+	if err != nil {
+		c.log.Warnf("Failed to update deployment %s: %s", deploymentName, err.Error())
+		return err
 	}
 	return nil
 }

rollout/sync_test.go

@@ -1,27 +1,35 @@
 package rollout
 
 import (
+	"fmt"
 	"strconv"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+
 	appsv1 "k8s.io/api/apps/v1"
 	v1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/util/intstr"
 	k8sfake "k8s.io/client-go/kubernetes/fake"
+	k8stesting "k8s.io/client-go/testing"
 	testclient "k8s.io/client-go/testing"
+	"k8s.io/client-go/tools/cache"
 	"k8s.io/utils/pointer"
 
 	"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
+
 	"github.com/argoproj/argo-rollouts/pkg/client/clientset/versioned/fake"
 	"github.com/argoproj/argo-rollouts/utils/annotations"
 	"github.com/argoproj/argo-rollouts/utils/conditions"
 	logutil "github.com/argoproj/argo-rollouts/utils/log"
 	"github.com/argoproj/argo-rollouts/utils/record"
 	timeutil "github.com/argoproj/argo-rollouts/utils/time"
+
+	"context"
 )
 
 func rs(name string, replicas int, selector map[string]string, timestamp metav1.Time, ownerRef *metav1.OwnerReference) *appsv1.ReplicaSet {
@@ -610,3 +618,230 @@ func Test_shouldFullPromote(t *testing.T) {
 	result = ctx.shouldFullPromote(newStatus)
 	assert.Equal(t, result, "Rollback within window")
 }
+
+type DeploymentActions interface {
+	Get(ctx context.Context, name string, opts metav1.GetOptions) (*appsv1.Deployment, error)
+	Update(ctx context.Context, deployment *appsv1.Deployment, opts metav1.UpdateOptions) (*appsv1.Deployment, error)
+}
+
+type KubeClientInterface interface {
+	Deployments(namespace string) DeploymentActions
+	AppsV1() AppV1Interface
+}
+
+type AppV1Interface interface {
+	Deployments(namespace string) DeploymentActions
+}
+
+type mockDeploymentInterface struct {
+	deployment *appsv1.Deployment
+}
+
+func (m *mockDeploymentInterface) Get(ctx context.Context, name string, opts metav1.GetOptions) (*appsv1.Deployment, error) {
+	return m.deployment, nil
+}
+
+func (m *mockDeploymentInterface) Update(ctx context.Context, deployment *appsv1.Deployment, opts metav1.UpdateOptions) (*appsv1.Deployment, error) {
+	m.deployment = deployment
+	return deployment, nil
+}
+
+type testKubeClient struct {
+	mockDeployment DeploymentActions
+}
+
+func (t *testKubeClient) AppsV1() AppV1Interface {
+	return t
+}
+
+func (t *testKubeClient) Deployments(namespace string) DeploymentActions {
+	return t.mockDeployment
+}
+
+type testRolloutContext struct {
+	*rolloutContext
+	kubeClient KubeClientInterface
+}
+
+func createTestRolloutContext(scaleDownMode string, initialReplicaSet *appsv1.ReplicaSet, deploymentReplicas int32, replicaSetInformer cache.SharedIndexInformer, deploymentExists bool, updateError error) *testRolloutContext {
+	ro := &v1alpha1.Rollout{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "rollout-test",
+			Namespace: "default",
+		},
+		Spec: v1alpha1.RolloutSpec{
+			WorkloadRef: &v1alpha1.ObjectRef{
+				Name:      "workload-test",
+				ScaleDown: scaleDownMode,
+			},
+		},
+	}
+
+	fakeDeployment := &appsv1.Deployment{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "workload-test",
+			Namespace: "default",
+		},
+		Spec: appsv1.DeploymentSpec{
+			Replicas: &deploymentReplicas,
+		},
+	}
+
+	var k8sfakeClient *k8sfake.Clientset
+	if deploymentExists {
+		if initialReplicaSet != nil {
+			k8sfakeClient = k8sfake.NewSimpleClientset(initialReplicaSet, fakeDeployment)
+		} else {
+			k8sfakeClient = k8sfake.NewSimpleClientset(fakeDeployment)
+		}
+	} else {
+		if initialReplicaSet != nil {
+			k8sfakeClient = k8sfake.NewSimpleClientset(initialReplicaSet)
+		} else {
+			k8sfakeClient = k8sfake.NewSimpleClientset()
+		}
+	}
+
+	if updateError != nil {
+		k8sfakeClient.PrependReactor("update", "deployments", func(action k8stesting.Action) (handled bool, ret runtime.Object, err error) {
+			return true, nil, updateError
+		})
+	}
+
+	mockDeploy := &mockDeploymentInterface{deployment: fakeDeployment}
+	testClient := &testKubeClient{mockDeployment: mockDeploy}
+
+	ctx := &testRolloutContext{
+		rolloutContext: &rolloutContext{
+			rollout:      ro,
+			pauseContext: &pauseContext{},
+			reconcilerBase: reconcilerBase{
+				argoprojclientset:  &fake.Clientset{},
+				kubeclientset:      k8sfakeClient,
+				recorder:           record.NewFakeEventRecorder(),
+				replicaSetInformer: replicaSetInformer,
+			},
+		},
+		kubeClient: testClient,
+	}
+	ctx.log = logutil.WithRollout(ctx.rollout)
+
+	return ctx
+}
+
+func TestScaleDownDeploymentOnSuccess(t *testing.T) {
+	ctx := createTestRolloutContext(v1alpha1.ScaleDownOnSuccess, nil, 5, nil, true, nil)
+	newStatus := &v1alpha1.RolloutStatus{
+		CurrentPodHash: "2f646bf702",
+		StableRS:       "15fb5ffc01",
+	}
+	err := ctx.promoteStable(newStatus, "reason")
+
+	assert.Nil(t, err)
+	k8sfakeClient := ctx.kubeclientset.(*k8sfake.Clientset)
+	updatedDeployment, err := k8sfakeClient.AppsV1().Deployments("default").Get(context.TODO(), "workload-test", metav1.GetOptions{})
+	assert.Nil(t, err)
+	assert.Equal(t, int32(0), *updatedDeployment.Spec.Replicas)
+
+	// test scale deployment error
+	ctx = createTestRolloutContext(v1alpha1.ScaleDownOnSuccess, nil, 5, nil, false, nil)
+	newStatus = &v1alpha1.RolloutStatus{
+		CurrentPodHash: "2f646bf702",
+		StableRS:       "15fb5ffc01",
+	}
+	err = ctx.promoteStable(newStatus, "reason")
+
+	assert.NotNil(t, err)
+}
+
+func TestScaleDownProgressively(t *testing.T) {
+	initialScale := int32(3)
+	initialReplicaSet := &appsv1.ReplicaSet{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "rs-test",
+			Namespace: "default",
+		},
+		Spec: appsv1.ReplicaSetSpec{
+			Replicas: &initialScale,
+		},
+	}
+
+	f := newFixture(t)
+	_, _, k8sInformer := f.newController(noResyncPeriodFunc)
+	replicaSetInformer := k8sInformer.Apps().V1().ReplicaSets().Informer()
+	ctx := createTestRolloutContext(v1alpha1.ScaleDownProgressively, initialReplicaSet, 5, replicaSetInformer, true, nil)
+	newScale := int32(5)
+	scaled, updatedRS, err := ctx.scaleReplicaSet(initialReplicaSet, newScale, ctx.rollout, "scaled up")
+
+	assert.Nil(t, err)
+	assert.True(t, scaled)
+	assert.Equal(t, newScale, *updatedRS.Spec.Replicas)
+	k8sfakeClient := ctx.kubeclientset.(*k8sfake.Clientset)
+	updatedDeployment, err := k8sfakeClient.AppsV1().Deployments("default").Get(context.TODO(), "workload-test", metav1.GetOptions{})
+	assert.Nil(t, err)
+	assert.Equal(t, int32(3), *updatedDeployment.Spec.Replicas)
+
+	// test scale deployment error
+	ctx = createTestRolloutContext(v1alpha1.ScaleDownProgressively, initialReplicaSet, 5, replicaSetInformer, false, nil)
+	_, _, err = ctx.scaleReplicaSet(initialReplicaSet, newScale, ctx.rollout, "scaled up")
+	assert.NotNil(t, err)
+}
+
+func TestScaleDeployment(t *testing.T) {
+	tests := []struct {
+		name             string
+		scaleToZero      bool
+		oldScale         *int32
+		newScale         *int32
+		expectedCount    int32
+		deploymentExists bool
+		updateError      error
+	}{
+		{
+			name:             "Scale down to zero",
+			scaleToZero:      true,
+			oldScale:         nil,
+			newScale:         nil,
+			expectedCount:    0,
+			deploymentExists: true,
+		},
+		{
+			name:             "Scale down by difference",
+			scaleToZero:      false,
+			oldScale:         int32Ptr(2),
+			newScale:         int32Ptr(3),
+			expectedCount:    4,
+			deploymentExists: true,
+		},
+		{
+			name:             "Error fetching deployment",
+			scaleToZero:      false,
+			oldScale:         int32Ptr(2),
+			newScale:         int32Ptr(3),
+			deploymentExists: false,
+		},
+		{
+			name:             "Error updating deployment",
+			scaleToZero:      false,
+			oldScale:         int32Ptr(2),
+			newScale:         int32Ptr(3),
+			deploymentExists: true,
+			updateError:      fmt.Errorf("fake update error"),
+		},
+	}
+
+	for _, test := range tests {
+		ctx := createTestRolloutContext(v1alpha1.ScaleDownOnSuccess, nil, 5, nil, test.deploymentExists, test.updateError)
+		err := ctx.scaleDeployment(test.scaleToZero, test.oldScale, test.newScale)
+
+		if !test.deploymentExists || test.updateError != nil {
+			assert.NotNil(t, err)
+			continue
+		}
+		assert.Nil(t, err)
+		k8sfakeClient := ctx.kubeclientset.(*k8sfake.Clientset)
+		updatedDeployment, err := k8sfakeClient.AppsV1().Deployments("default").Get(context.TODO(), "workload-test", metav1.GetOptions{})
+		assert.Nil(t, err)
+		assert.Equal(t, *updatedDeployment.Spec.Replicas, test.expectedCount)
+	}
+}

ui/src/models/rollout/generated/api.ts

@@ -1700,6 +1700,12 @@ export interface GithubComArgoprojArgoRolloutsPkgApisRolloutsV1alpha1ObjectRef {
      * @memberof GithubComArgoprojArgoRolloutsPkgApisRolloutsV1alpha1ObjectRef
      */
     name?: string;
+    /**
+     * 
+     * @type {string}
+     * @memberof GithubComArgoprojArgoRolloutsPkgApisRolloutsV1alpha1ObjectRef
+     */
+    scaleDown?: string;
 }
 /**
  *