Feat: support rollout controller for StatefulSet (#1969)

* Feat: support rollout controller for StatefulSet

* Feat: support one statefulset in rollout

* Feat: add tests for StatefulSet rollout controller

* Fix: correct workload-deleting error

* Feat: remove advanced sts

pkg/controller/common/rollout/rollout_plan_controller.go

@@ -369,6 +369,17 @@ func (r *Controller) GetWorkloadController() (workloads.WorkloadController, erro
 			return workloads.NewDeploymentScaleController(r.client, r.recorder, r.parentController,
 				r.rolloutSpec, r.rolloutStatus, target), nil
 		}
+
+		// check if the target workload is StatefulSet
+		if r.targetWorkload.GetKind() == reflect.TypeOf(apps.StatefulSet{}).Name() {
+			// check whether current rollout plan is for workload rolling or scaling
+			if r.sourceWorkload != nil {
+				return workloads.NewStatefulSetRolloutController(r.client, r.recorder, r.parentController,
+					r.rolloutSpec, r.rolloutStatus, target), nil
+			}
+			return workloads.NewStatefulSetScaleController(r.client, r.recorder, r.parentController,
+				r.rolloutSpec, r.rolloutStatus, target), nil
+		}
 	}
 
 	return nil, fmt.Errorf("the workload kind `%s` is not supported", kind)

pkg/controller/common/rollout/workloads/statefulset_controller.go

@@ -29,6 +29,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	"github.com/oam-dev/kubevela/apis/core.oam.dev/v1beta1"
+	"github.com/oam-dev/kubevela/apis/standard.oam.dev/v1alpha1"
 )
 
 type statefulSetController struct {
@@ -38,9 +39,10 @@ type statefulSetController struct {
 
 // add the parent controller to the owner of the StatefulSet, and initialize the size
 // before kicking start the update and start from every pod in the old version
-func (c *statefulSetController) claimStatefulSet(ctx context.Context, statefulSet *apps.StatefulSet, initSize *int32) (bool, error) {
+func (c *statefulSetController) claimStatefulSet(ctx context.Context, statefulSet *apps.StatefulSet) (bool, error) {
 	if controller := metav1.GetControllerOf(statefulSet); controller != nil &&
-		controller.Kind == v1beta1.AppRolloutKind && controller.APIVersion == v1beta1.SchemeGroupVersion.String() {
+		(controller.Kind == v1beta1.AppRolloutKind && controller.APIVersion == v1beta1.SchemeGroupVersion.String() ||
+			controller.Kind == v1alpha1.RolloutKind && controller.APIVersion == v1alpha1.SchemeGroupVersion.String()) {
 		// it's already there
 		return true, nil
 	}
@@ -48,13 +50,9 @@ func (c *statefulSetController) claimStatefulSet(ctx context.Context, statefulSe
 	statefulSetPatch := client.MergeFrom(statefulSet.DeepCopy())
 
 	// add the parent controller to the owner of the StatefulSet
-	ref := metav1.NewControllerRef(c.parentController, v1beta1.AppRolloutKindVersionKind)
+	ref := metav1.NewControllerRef(c.parentController, c.parentController.GetObjectKind().GroupVersionKind())
 	statefulSet.SetOwnerReferences(append(statefulSet.GetOwnerReferences(), *ref))
 
-	if initSize != nil {
-		statefulSet.Spec.Replicas = initSize
-	}
-
 	// patch the StatefulSet
 	if err := c.client.Patch(ctx, statefulSet, statefulSetPatch, client.FieldOwner(c.parentController.GetUID())); err != nil {
 		c.recorder.Event(c.parentController, event.Warning("Failed to the start the StatefulSet update", err))
@@ -82,6 +80,23 @@ func (c *statefulSetController) scaleStatefulSet(ctx context.Context, statefulSe
 	return nil
 }
 
+func (c *statefulSetController) setPartition(ctx context.Context, statefulSet *apps.StatefulSet, partition int32) error {
+	statefulSetPatch := client.MergeFrom(statefulSet.DeepCopy())
+	statefulSet.Spec.UpdateStrategy.RollingUpdate.Partition = pointer.Int32Ptr(partition)
+
+	// patch the StatefulSet
+	if err := c.client.Patch(ctx, statefulSet, statefulSetPatch, client.FieldOwner(c.parentController.GetUID())); err != nil {
+		c.recorder.Event(c.parentController, event.Warning(event.Reason(fmt.Sprintf(
+			"Failed to update the partition of StatefulSet %s to the correct target %d", statefulSet.GetName(), partition)), err))
+		c.rolloutStatus.RolloutRetry(err.Error())
+		return err
+	}
+
+	klog.InfoS("Submitted upgrade quest for StatefulSet", "StatefulSet",
+		statefulSet.GetName(), "target partition", partition, "batch", c.rolloutStatus.CurrentBatch)
+	return nil
+}
+
 // remove the parent controller from the StatefulSet's owner list
 func (c *statefulSetController) releaseStatefulSet(ctx context.Context, statefulSet *apps.StatefulSet) (bool, error) {
 	statefulSetPatch := client.MergeFrom(statefulSet.DeepCopy())

pkg/controller/common/rollout/workloads/statefulset_rollout_controller.go

@@ -0,0 +1,291 @@
+/*
+Copyright 2021 The KubeVela Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package workloads
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/crossplane/crossplane-runtime/pkg/event"
+	"github.com/pkg/errors"
+	appsv1 "k8s.io/api/apps/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/intstr"
+	"k8s.io/klog/v2"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	"github.com/oam-dev/kubevela/apis/standard.oam.dev/v1alpha1"
+	"github.com/oam-dev/kubevela/pkg/oam"
+)
+
+// StatefulSetRolloutController is responsible for handle rollout StatefulSet type of workloads
+type StatefulSetRolloutController struct {
+	statefulSetController
+	statefulSet *appsv1.StatefulSet
+}
+
+// NewStatefulSetRolloutController creates StatefulSet rollout controller
+func NewStatefulSetRolloutController(client client.Client, recorder event.Recorder, parentController oam.Object, rolloutSpec *v1alpha1.RolloutPlan, rolloutStatus *v1alpha1.RolloutStatus,
+	targetNamespacedName types.NamespacedName) *StatefulSetRolloutController {
+	return &StatefulSetRolloutController{
+		statefulSetController: statefulSetController{
+			workloadController: workloadController{
+				client:           client,
+				recorder:         recorder,
+				parentController: parentController,
+				rolloutSpec:      rolloutSpec,
+				rolloutStatus:    rolloutStatus,
+			},
+			targetNamespacedName: targetNamespacedName,
+		},
+	}
+}
+
+// VerifySpec verifies that the rollout resource is consistent with the rollout spec
+func (s *StatefulSetRolloutController) VerifySpec(ctx context.Context) (bool, error) {
+	var verifyErr error
+
+	defer func() {
+		if verifyErr != nil {
+			klog.Error(verifyErr)
+			s.recorder.Event(s.parentController, event.Warning("VerifyFailed", verifyErr))
+		}
+	}()
+
+	currentReplicas, verifyErr := s.size(ctx)
+	if verifyErr != nil {
+		s.rolloutStatus.RolloutRetry(verifyErr.Error())
+		// nolint: nilerr
+		return false, nil
+	}
+	// record the size and we will use this value to drive the rest of the batches
+	klog.InfoS("record the target size", "total replicas", currentReplicas)
+	s.rolloutStatus.RolloutTargetSize = currentReplicas
+	s.rolloutStatus.RolloutOriginalSize = currentReplicas
+
+	// make sure that the updateRevision is different from what we have already done
+	targetHash := s.statefulSet.Status.UpdateRevision
+	if targetHash == s.rolloutStatus.LastAppliedPodTemplateIdentifier {
+		return false, fmt.Errorf("there is no difference between the source and target, hash = %s", targetHash)
+	}
+
+	if s.statefulSet.Spec.Replicas != nil && currentReplicas != s.statefulSet.Status.Replicas {
+		verifyErr = fmt.Errorf("the StatefulSet is still scaling, target = %d, statefulSet size = %d",
+			currentReplicas, s.statefulSet.Status.Replicas)
+		s.rolloutStatus.RolloutRetry(verifyErr.Error())
+		return false, verifyErr
+	}
+
+	// check if the rollout batch replicas added up to the StatefulSet replicas
+	if verifyErr = s.verifyRolloutBatchReplicaValue(currentReplicas); verifyErr != nil {
+		return false, verifyErr
+	}
+
+	// check if the StatefulSet has any controller
+	if controller := metav1.GetControllerOf(s.statefulSet); controller != nil {
+		return false, fmt.Errorf("the StatefulSet %s has a controller owner %s",
+			s.statefulSet.GetName(), controller.String())
+	}
+
+	// mark the rollout verified
+	s.recorder.Event(s.parentController, event.Normal("Rollout Verified",
+		"Rollout spec and the StatefulSet resource are verified"))
+	// record the new pod template StatefulSet on success
+	s.rolloutStatus.NewPodTemplateIdentifier = targetHash
+	return true, nil
+}
+
+// Initialize makes sure that the source and target StatefulSet is under our control
+func (s *StatefulSetRolloutController) Initialize(ctx context.Context) (bool, error) {
+	currentReplicas, err := s.size(ctx)
+	if err != nil {
+		s.rolloutStatus.RolloutRetry(err.Error())
+		return false, nil
+	}
+
+	if _, err := s.claimStatefulSet(ctx, s.statefulSet); err != nil {
+		// nolint:nilerr
+		return false, nil
+	}
+
+	if err := s.setPartition(ctx, s.statefulSet, currentReplicas); err != nil {
+		// nolint:nilerr
+		return false, nil
+	}
+
+	// mark the rollout initialized
+	s.recorder.Event(s.parentController, event.Normal("Rollout Initialized", "Rollout resource are initialized"))
+	return true, nil
+}
+
+// RolloutOneBatchPods calculates the number of pods we can upgrade once according to the rollout spec
+// and then set the partition accordingly
+func (s *StatefulSetRolloutController) RolloutOneBatchPods(ctx context.Context) (bool, error) {
+	currentReplicas, err := s.size(ctx)
+	if err != nil {
+		s.rolloutStatus.RolloutRetry(err.Error())
+		return false, nil
+	}
+
+	newPodTarget := s.calculateCurrentTarget(currentReplicas)
+	if err := s.setPartition(ctx, s.statefulSet, currentReplicas-newPodTarget); err != nil {
+		// nolint:nilerr
+		return false, nil
+	}
+
+	// record the finished upgrade action
+	klog.InfoS("upgraded one batch", "current batch", s.rolloutStatus.CurrentBatch,
+		"target size", newPodTarget)
+	s.recorder.Event(s.parentController, event.Normal("Batch Rollout",
+		fmt.Sprintf("Finished submiting all upgrade quests for batch %d", s.rolloutStatus.CurrentBatch)))
+	s.rolloutStatus.UpgradedReplicas = newPodTarget
+	return true, nil
+}
+
+// CheckOneBatchPods checks to see if the pods are all available according to the rollout plan
+func (s *StatefulSetRolloutController) CheckOneBatchPods(ctx context.Context) (bool, error) {
+	currentReplicas, err := s.size(ctx)
+	if err != nil {
+		s.rolloutStatus.RolloutRetry(err.Error())
+		return false, nil
+	}
+
+	newPodTarget := s.calculateCurrentTarget(currentReplicas)
+	readyPodCount := int(s.statefulSet.Status.ReadyReplicas)
+
+	if len(s.rolloutSpec.RolloutBatches) <= int(s.rolloutStatus.CurrentBatch) {
+		err := errors.New("somehow, currentBatch number exceeded the rolloutBatches spec")
+		klog.ErrorS(err, "total batch", len(s.rolloutSpec.RolloutBatches), "current batch",
+			s.rolloutStatus.CurrentBatch)
+		return false, err
+	}
+
+	currentBatch := s.rolloutSpec.RolloutBatches[s.rolloutStatus.CurrentBatch]
+	maxUnavail := 0
+	if currentBatch.MaxUnavailable != nil {
+		maxUnavail, _ = intstr.GetValueFromIntOrPercent(currentBatch.MaxUnavailable, int(currentReplicas), true)
+	}
+	klog.InfoS("checking the rolling out progress", "current batch", s.rolloutStatus.CurrentBatch,
+		"new pod count target", newPodTarget, "new ready pod count", readyPodCount,
+		"max unavailable pod allowed", maxUnavail)
+	s.rolloutStatus.UpgradedReadyReplicas = int32(readyPodCount)
+
+	if maxUnavail+readyPodCount >= int(newPodTarget) {
+		// record the successful upgrade
+		klog.InfoS("all pods in current batch are ready", "current batch", s.rolloutStatus.CurrentBatch)
+		s.recorder.Event(s.parentController, event.Normal("Batch Available",
+			fmt.Sprintf("Batch %d is available", s.rolloutStatus.CurrentBatch)))
+		return true, nil
+	}
+
+	// continue to verify
+	klog.InfoS("the batch is not ready yet", "current batch", s.rolloutStatus.CurrentBatch)
+	s.rolloutStatus.RolloutRetry("the batch is not ready yet")
+	return false, nil
+}
+
+// FinalizeOneBatch makes sure that the rollout status are updated correctly
+func (s *StatefulSetRolloutController) FinalizeOneBatch(ctx context.Context) (bool, error) {
+	status := s.rolloutStatus
+	spec := s.rolloutSpec
+
+	if spec.BatchPartition != nil && *spec.BatchPartition < status.CurrentBatch {
+		err := fmt.Errorf("the current batch value in the status is greater than the batch partition")
+		klog.ErrorS(err, "we have moved past the user defined partition", "user specified batch partition",
+			*spec.BatchPartition, "current batch we are working on", status.CurrentBatch)
+		return false, err
+	}
+
+	upgradedReplicas := int(status.UpgradedReplicas)
+	currentBatch := int(status.CurrentBatch)
+	// calculate the lower bound of the possible pod count just before the current batch
+	podCount := calculateNewBatchTarget(s.rolloutSpec, 0, int(s.rolloutStatus.RolloutTargetSize), currentBatch-1)
+	// the recorded number should be at least as much as the all the pods before the current batch
+	if podCount > upgradedReplicas {
+		err := fmt.Errorf("the upgraded replica in the status is less than all the pods in the previous batch")
+		klog.ErrorS(err, "rollout status inconsistent", "upgraded num status", upgradedReplicas,
+			"pods in all the previous batches", podCount)
+		return false, err
+	}
+
+	// calculate the upper bound with the current batch
+	podCount = calculateNewBatchTarget(s.rolloutSpec, 0, int(s.rolloutStatus.RolloutTargetSize), currentBatch)
+	// the recorded number should be not as much as the all the pods including the active batch
+	if podCount < upgradedReplicas {
+		err := fmt.Errorf("the upgraded replica in the status is greater than all the pods in the current batch")
+		klog.ErrorS(err, "rollout status inconsistent", "total target size", s.rolloutStatus.RolloutTargetSize,
+			"upgraded num status", upgradedReplicas, "pods in the batches including the current batch", podCount)
+		return false, err
+	}
+	return true, nil
+}
+
+// Finalize makes sure the StatefulSet is all upgraded
+func (s *StatefulSetRolloutController) Finalize(ctx context.Context, succeed bool) bool {
+	if err := s.fetchStatefulSet(ctx); err != nil {
+		// don't fail the rollout just because of we can't get the resource
+		return false
+	}
+
+	// release StatefulSet
+	if _, err := s.releaseStatefulSet(ctx, s.statefulSet); err != nil {
+		return false
+	}
+
+	// mark the resource finalized
+	s.rolloutStatus.LastAppliedPodTemplateIdentifier = s.rolloutStatus.NewPodTemplateIdentifier
+	s.recorder.Event(s.parentController, event.Normal("Rollout Finalized",
+		fmt.Sprintf("Rollout resource are finalized, succeed := %t", succeed)))
+	return true
+}
+
+// check if the replicas in all the rollout batches add up to the right number
+func (s *StatefulSetRolloutController) verifyRolloutBatchReplicaValue(totalReplicas int32) error {
+	return verifyBatchesWithRollout(s.rolloutSpec, totalReplicas)
+}
+
+// the target StatefulSet size for the current batch
+func (s *StatefulSetRolloutController) calculateCurrentTarget(totalSize int32) int32 {
+	targetSize := int32(calculateNewBatchTarget(s.rolloutSpec, 0, int(totalSize), int(s.rolloutStatus.CurrentBatch)))
+	klog.InfoS("Calculated the number of pods in the target StatefulSet after current batch",
+		"current batch", s.rolloutStatus.CurrentBatch, "target StatefulSet size", targetSize)
+	return targetSize
+}
+
+func (s *StatefulSetRolloutController) fetchStatefulSet(ctx context.Context) error {
+	workload := appsv1.StatefulSet{}
+	if err := s.client.Get(ctx, s.targetNamespacedName, &workload); err != nil {
+		if !apierrors.IsNotFound(err) {
+			s.recorder.Event(s.parentController, event.Warning("Failed to get the StatefulSet", err))
+		}
+		return err
+	}
+	s.statefulSet = &workload
+	return nil
+}
+
+func (s *StatefulSetRolloutController) size(ctx context.Context) (int32, error) {
+	if s.statefulSet == nil {
+		if err := s.fetchStatefulSet(ctx); err != nil {
+			return 0, err
+		}
+	}
+	// default is 1
+	return getStatefulSetReplicas(s.statefulSet), nil
+}

pkg/controller/common/rollout/workloads/statefulset_rollout_integration_test.go

@@ -0,0 +1,500 @@
+/*
+
+ Copyright 2021 The KubeVela Authors.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+
+*/
+
+package workloads
+
+import (
+	"fmt"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+
+	"github.com/crossplane/crossplane-runtime/pkg/event"
+	apps "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/intstr"
+	"k8s.io/utils/pointer"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	"github.com/oam-dev/kubevela/apis/core.oam.dev/v1beta1"
+	"github.com/oam-dev/kubevela/apis/standard.oam.dev/v1alpha1"
+	"github.com/oam-dev/kubevela/pkg/oam/util"
+)
+
+var _ = Describe("StatefulSet controller", func() {
+	var (
+		c              StatefulSetRolloutController
+		ns             corev1.Namespace
+		name           string
+		namespace      string
+		statefulSet    apps.StatefulSet
+		namespacedName client.ObjectKey
+	)
+
+	BeforeEach(func() {
+		namespace = "rollout-ns"
+		name = "rollout"
+		appRollout := v1beta1.AppRollout{TypeMeta: metav1.TypeMeta{APIVersion: v1beta1.SchemeGroupVersion.String(), Kind: v1beta1.AppRolloutKind}, ObjectMeta: metav1.ObjectMeta{Name: name}}
+		namespacedName = client.ObjectKey{Name: name, Namespace: namespace}
+		c = StatefulSetRolloutController{
+			statefulSetController: statefulSetController{
+				workloadController: workloadController{
+					client: k8sClient,
+					rolloutSpec: &v1alpha1.RolloutPlan{
+						RolloutBatches: []v1alpha1.RolloutBatch{
+							{
+								Replicas: intstr.FromInt(1),
+							},
+						},
+					},
+					rolloutStatus:    &v1alpha1.RolloutStatus{RollingState: v1alpha1.RolloutSucceedState},
+					parentController: &appRollout,
+					recorder: event.NewAPIRecorder(mgr.GetEventRecorderFor("AppRollout")).
+						WithAnnotations("controller", "AppRollout"),
+				},
+				targetNamespacedName: namespacedName,
+			},
+		}
+
+		statefulSet = apps.StatefulSet{
+			TypeMeta:   metav1.TypeMeta{APIVersion: apps.SchemeGroupVersion.String(), Kind: "StatefulSet"},
+			ObjectMeta: metav1.ObjectMeta{Namespace: namespace, Name: name},
+			Spec: apps.StatefulSetSpec{
+				Selector: &metav1.LabelSelector{
+					MatchLabels: map[string]string{"env": "staging"},
+				},
+				Template: corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{"env": "staging"}},
+					Spec:       corev1.PodSpec{Containers: []corev1.Container{{Name: name, Image: "nginx"}}},
+				},
+			},
+		}
+
+		ns = corev1.Namespace{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: namespace,
+			},
+		}
+		By("Create a namespace")
+		Expect(k8sClient.Create(ctx, &ns)).Should(SatisfyAny(Succeed(), &util.AlreadyExistMatcher{}))
+	})
+
+	AfterEach(func() {
+		By("clean up")
+		k8sClient.Delete(ctx, &statefulSet)
+	})
+
+	Context("TestNewStatefulSetRolloutController", func() {
+		It("init a StatefulSet Rollout Controller", func() {
+			recorder := event.NewAPIRecorder(mgr.GetEventRecorderFor("AppRollout")).
+				WithAnnotations("controller", "AppRollout")
+			parentController := &v1beta1.AppRollout{ObjectMeta: metav1.ObjectMeta{Name: name}}
+			rolloutSpec := &v1alpha1.RolloutPlan{
+				RolloutBatches: []v1alpha1.RolloutBatch{{
+					Replicas: intstr.FromInt(1),
+				},
+				},
+			}
+			rolloutStatus := &v1alpha1.RolloutStatus{RollingState: v1alpha1.RolloutSucceedState}
+			workloadNamespacedName := client.ObjectKey{Name: name, Namespace: namespace}
+			got := NewStatefulSetRolloutController(k8sClient, recorder, parentController, rolloutSpec, rolloutStatus, workloadNamespacedName)
+			c := &StatefulSetRolloutController{
+				statefulSetController: statefulSetController{
+					workloadController: workloadController{
+						client:           k8sClient,
+						recorder:         recorder,
+						parentController: parentController,
+						rolloutSpec:      rolloutSpec,
+						rolloutStatus:    rolloutStatus,
+					},
+					targetNamespacedName: workloadNamespacedName,
+				},
+			}
+			Expect(got).Should(Equal(c))
+		})
+	})
+
+	Context("VerifySpec", func() {
+		It("could not fetch StatefulSet workload", func() {
+			consistent, err := c.VerifySpec(ctx)
+			Expect(err).Should(BeNil())
+			Expect(consistent).Should(BeFalse())
+		})
+
+		It("verify rollout spec hash", func() {
+			By("Create a StatefulSet")
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("Verify should fail because the the target hash didn't change")
+			targetHash := statefulSet.Status.UpdateRevision
+			c.rolloutStatus.LastAppliedPodTemplateIdentifier = targetHash
+			consistent, err := c.VerifySpec(ctx)
+			Expect(err).Should(Equal(fmt.Errorf("there is no difference between the source and target, hash = ")))
+			Expect(consistent).Should(BeFalse())
+		})
+
+		It("the StatefulSet need to be stable", func() {
+			By("create the StatefulSet with many pods")
+			statefulSet.Spec.Replicas = pointer.Int32Ptr(50)
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("setting a dummy pod identifier so it's different")
+			c.rolloutStatus.LastAppliedPodTemplateIdentifier = "abc"
+
+			By("verify should fail because the StatefulSet is not stable")
+			consistent, err := c.VerifySpec(ctx)
+			Expect(consistent).Should(BeFalse())
+			Expect(err.Error()).Should(ContainSubstring("is still scaling"))
+			Expect(c.rolloutStatus.RolloutTargetSize).Should(BeEquivalentTo(50))
+			Expect(c.rolloutStatus.NewPodTemplateIdentifier).Should(BeEmpty())
+		})
+
+		It("the StatefulSet should not have controller", func() {
+			By("Create a StatefulSet")
+			statefulSet.SetOwnerReferences([]metav1.OwnerReference{{
+				APIVersion: v1beta1.SchemeGroupVersion.String(),
+				Kind:       v1beta1.ApplicationKind,
+				Name:       "def",
+				UID:        "123456",
+				Controller: pointer.BoolPtr(true),
+			}})
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("setting a dummy pod identifier so it's different")
+			c.rolloutStatus.LastAppliedPodTemplateIdentifier = "abc"
+
+			statefulSet.Status.Replicas = *statefulSet.Spec.Replicas
+			Expect(k8sClient.Status().Update(ctx, &statefulSet)).Should(Succeed())
+
+			By("verify should fail because the StatefulSet still has a controller")
+			consistent, err := c.VerifySpec(ctx)
+			Expect(consistent).Should(BeFalse())
+			Expect(err.Error()).Should(ContainSubstring("has a controller owner"))
+		})
+
+		It("spec is valid", func() {
+			By("Create a StatefulSet")
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("setting a dummy pod identifier so it's different")
+			c.rolloutStatus.LastAppliedPodTemplateIdentifier = "abc"
+
+			statefulSet.Status.Replicas = *statefulSet.Spec.Replicas
+			Expect(k8sClient.Status().Update(ctx, &statefulSet)).Should(Succeed())
+
+			By("verify should succeed")
+			consistent, err := c.VerifySpec(ctx)
+			Expect(err).Should(BeNil())
+			Expect(consistent).Should(BeTrue())
+			Expect(c.rolloutStatus.RolloutTargetSize).Should(BeEquivalentTo(*statefulSet.Spec.Replicas))
+			Expect(c.rolloutStatus.NewPodTemplateIdentifier).Should(BeEmpty())
+		})
+	})
+
+	Context("TestInitialize", func() {
+		It("could not fetch StatefulSet workload", func() {
+			consistent, err := c.Initialize(ctx)
+			Expect(err).Should(BeNil())
+			Expect(consistent).Should(BeFalse())
+		})
+
+		It("failed to patch the owner of StatefulSet", func() {
+			By("Create a StatefulSet")
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("initialize will fail because StatefulSet has wrong owner reference")
+			initialized, err := c.Initialize(ctx)
+			Expect(initialized).Should(BeFalse())
+			Expect(err).Should(BeNil())
+		})
+
+		It("workload StatefulSet is controlled by appRollout already", func() {
+			By("Create a StatefulSet")
+			statefulSet.SetOwnerReferences([]metav1.OwnerReference{{
+				APIVersion: v1beta1.SchemeGroupVersion.String(),
+				Kind:       v1beta1.AppRolloutKind,
+				Name:       "def",
+				UID:        "123456",
+				Controller: pointer.BoolPtr(true),
+			}})
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("initialize succeed without patching")
+			initialized, err := c.Initialize(ctx)
+			Expect(initialized).Should(BeTrue())
+			Expect(err).Should(BeNil())
+			Expect(k8sClient.Get(ctx, c.targetNamespacedName, &statefulSet)).Should(Succeed())
+			Expect(len(statefulSet.GetOwnerReferences())).Should(BeEquivalentTo(1))
+		})
+
+		It("successfully initialized StatefulSet", func() {
+			By("create StatefulSet")
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("initialize succeeds")
+			c.parentController.SetUID("1231586900")
+			initialized, err := c.Initialize(ctx)
+			Expect(initialized).Should(BeTrue())
+			Expect(err).Should(BeNil())
+			Expect(k8sClient.Get(ctx, c.targetNamespacedName, &statefulSet)).Should(Succeed())
+			Expect(len(statefulSet.GetOwnerReferences())).Should(BeEquivalentTo(1))
+		})
+	})
+
+	Context("TestRolloutOneBatchPods", func() {
+		It("could not fetch StatefulSet workload", func() {
+			consistent, err := c.RolloutOneBatchPods(ctx)
+			Expect(err).Should(BeNil())
+			Expect(consistent).Should(BeFalse())
+		})
+
+		It("successfully rollout, current batch number is not equal to the expected one", func() {
+			By("Create a StatefulSet")
+			statefulSet.Spec.Replicas = pointer.Int32Ptr(10)
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("rollout the second batch of current StatefulSet")
+			c.rolloutStatus.CurrentBatch = 1
+			c.rolloutSpec.RolloutBatches = []v1alpha1.RolloutBatch{
+				{
+					Replicas: intstr.FromInt(1),
+				},
+				{
+					Replicas: intstr.FromString("20%"),
+				},
+				{
+					Replicas: intstr.FromString("80%"),
+				},
+			}
+			done, err := c.RolloutOneBatchPods(ctx)
+			Expect(done).Should(BeTrue())
+			Expect(err).Should(BeNil())
+			Expect(c.rolloutStatus.UpgradedReplicas).Should(BeEquivalentTo(3))
+			Expect(k8sClient.Get(ctx, c.targetNamespacedName, &statefulSet)).Should(Succeed())
+			Expect(*statefulSet.Spec.UpdateStrategy.RollingUpdate.Partition).Should(BeEquivalentTo(7))
+		})
+	})
+
+	Context("TestCheckOneBatchPods", func() {
+		BeforeEach(func() {
+			statefulSet.Spec.Replicas = pointer.Int32Ptr(10)
+			c.rolloutSpec.RolloutBatches = []v1alpha1.RolloutBatch{
+				{
+					Replicas: intstr.FromInt(2),
+				},
+				{
+					Replicas: intstr.FromString("20%"),
+				},
+				{
+					Replicas: intstr.FromString("80%"),
+				},
+			}
+		})
+
+		It("could not fetch StatefulSet workload", func() {
+			consistent, err := c.CheckOneBatchPods(ctx)
+			Expect(err).Should(BeNil())
+			Expect(consistent).Should(BeFalse())
+		})
+
+		It("current ready Pod is less than expected", func() {
+			By("Create the StatefulSet")
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+			By("Update the StatefulSet status")
+			statefulSet.Status.Replicas = 4
+			statefulSet.Status.ReadyReplicas = 3
+			statefulSet.Status.UpdatedReplicas = 4
+			Expect(k8sClient.Status().Update(ctx, &statefulSet)).Should(Succeed())
+
+			By("checking should fail as not enough pod ready")
+			c.rolloutStatus.CurrentBatch = 1
+			done, err := c.CheckOneBatchPods(ctx)
+			Expect(done).Should(BeFalse())
+			Expect(err).Should(BeNil())
+			Expect(c.rolloutStatus.UpgradedReadyReplicas).Should(BeEquivalentTo(statefulSet.Status.ReadyReplicas))
+		})
+
+		It("failed to check batch Pod when current batch number exceeds the expected ones", func() {
+			By("Create a StatefulSet")
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("checking")
+			c.rolloutStatus.CurrentBatch = 3
+			done, err := c.CheckOneBatchPods(ctx)
+			Expect(done).Should(BeFalse())
+			Expect(err.Error()).Should(ContainSubstring("currentBatch number exceeded the rolloutBatches spec"))
+		})
+
+		It("there are enough pods counting the unavailable", func() {
+			By("Create the StatefulSet")
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+			By("Update the StatefulSet status")
+			statefulSet.Status.Replicas = 4
+			statefulSet.Status.ReadyReplicas = 3
+			statefulSet.Status.UpdatedReplicas = 4
+			Expect(k8sClient.Status().Update(ctx, &statefulSet)).Should(Succeed())
+			c.rolloutStatus.CurrentBatch = 1
+			// set the rollout batch spec allow unavailable
+			perc := intstr.FromString("20%")
+			c.rolloutSpec.RolloutBatches = []v1alpha1.RolloutBatch{
+				{
+					Replicas: intstr.FromInt(2),
+				},
+				{
+					Replicas:       perc,
+					MaxUnavailable: &perc,
+				},
+				{
+					Replicas: intstr.FromString("80%"),
+				},
+			}
+			By("checking one batch")
+			done, err := c.CheckOneBatchPods(ctx)
+			Expect(done).Should(BeTrue())
+			Expect(err).Should(BeNil())
+			Expect(c.rolloutStatus.UpgradedReadyReplicas).Should(BeEquivalentTo(statefulSet.Status.ReadyReplicas))
+		})
+
+		It("there are enough pods ready", func() {
+			By("Create the StatefulSet")
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+			By("Update the StatefulSet status")
+			statefulSet.Status.Replicas = 10
+			statefulSet.Status.ReadyReplicas = 10
+			statefulSet.Status.UpdatedReplicas = 10
+			Expect(k8sClient.Status().Update(ctx, &statefulSet)).Should(Succeed())
+
+			By("the second batch should pass when there are more pods upgraded already")
+			c.rolloutStatus.CurrentBatch = 1
+			done, err := c.CheckOneBatchPods(ctx)
+			Expect(done).Should(BeTrue())
+			Expect(err).Should(BeNil())
+			Expect(c.rolloutStatus.UpgradedReadyReplicas).Should(BeEquivalentTo(statefulSet.Status.ReadyReplicas))
+
+			By("checking the last batch")
+			c.rolloutStatus.CurrentBatch = 2
+			done, err = c.CheckOneBatchPods(ctx)
+			Expect(done).Should(BeTrue())
+			Expect(err).Should(BeNil())
+			Expect(c.rolloutStatus.UpgradedReadyReplicas).Should(BeEquivalentTo(statefulSet.Status.ReadyReplicas))
+		})
+	})
+
+	Context("TestFinalizeOneBatch", func() {
+		BeforeEach(func() {
+			c.rolloutStatus.RolloutTargetSize = 10
+			c.rolloutSpec.RolloutBatches = []v1alpha1.RolloutBatch{
+				{
+					Replicas: intstr.FromInt(2),
+				},
+				{
+					Replicas: intstr.FromString("20%"),
+				},
+				{
+					Replicas: intstr.FromString("80%"),
+				},
+			}
+		})
+
+		It("test illegal batch partition", func() {
+			By("finalizing one batch")
+			c.rolloutSpec.BatchPartition = pointer.Int32Ptr(2)
+			c.rolloutStatus.CurrentBatch = 3
+			done, err := c.FinalizeOneBatch(ctx)
+			Expect(done).Should(BeFalse())
+			Expect(err.Error()).Should(ContainSubstring("the current batch value in the status is greater than the batch partition"))
+		})
+
+		It("test too few upgraded", func() {
+			By("finalizing one batch")
+			c.rolloutStatus.UpgradedReplicas = 2
+			c.rolloutStatus.CurrentBatch = 2
+			done, err := c.FinalizeOneBatch(ctx)
+			Expect(done).Should(BeFalse())
+			Expect(err.Error()).Should(ContainSubstring("is less than all the pods in the previous batch"))
+		})
+
+		It("test too many upgraded", func() {
+			By("finalizing one batch")
+			c.rolloutStatus.UpgradedReplicas = 5
+			c.rolloutStatus.CurrentBatch = 1
+			done, err := c.FinalizeOneBatch(ctx)
+			Expect(done).Should(BeFalse())
+			Expect(err.Error()).Should(ContainSubstring("is greater than all the pods in the current batch"))
+		})
+
+		It("test upgraded in the range", func() {
+			By("finalizing one batch")
+			c.rolloutStatus.UpgradedReplicas = 3
+			c.rolloutStatus.CurrentBatch = 1
+			done, err := c.FinalizeOneBatch(ctx)
+			Expect(done).Should(BeTrue())
+			Expect(err).Should(BeNil())
+		})
+	})
+
+	Context("TestFinalize", func() {
+		It("failed to fetch StatefulSet", func() {
+			By("finalizing")
+			finalized := c.Finalize(ctx, true)
+			Expect(finalized).Should(BeFalse())
+		})
+
+		It("Already finalize StatefulSet", func() {
+			By("Create a StatefulSet")
+			statefulSet.SetOwnerReferences([]metav1.OwnerReference{{
+				APIVersion: v1beta1.SchemeGroupVersion.String(),
+				Kind:       "notRollout",
+				Name:       "def",
+				UID:        "123456",
+			}})
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("finalizing without patch")
+			finalized := c.Finalize(ctx, true)
+			Expect(finalized).Should(BeTrue())
+		})
+
+		It("successfully to finalize StatefulSet", func() {
+			By("Create a StatefulSet")
+			statefulSet.SetOwnerReferences([]metav1.OwnerReference{
+				{
+					APIVersion: v1beta1.SchemeGroupVersion.String(),
+					Kind:       v1beta1.AppRolloutKind,
+					Name:       "def",
+					UID:        "123456",
+				},
+				{
+					APIVersion: corev1.SchemeGroupVersion.String(),
+					Kind:       "Deployment",
+					Name:       "def",
+					UID:        "998877745",
+				},
+			})
+			Expect(k8sClient.Create(ctx, &statefulSet)).Should(Succeed())
+
+			By("finalizing with patch")
+			finalized := c.Finalize(ctx, false)
+			Expect(finalized).Should(BeTrue())
+			Expect(k8sClient.Get(ctx, c.targetNamespacedName, &statefulSet)).Should(Succeed())
+			Expect(len(statefulSet.GetOwnerReferences())).Should(BeEquivalentTo(1))
+			Expect(statefulSet.GetOwnerReferences()[0].Kind).Should(Equal("Deployment"))
+		})
+	})
+})

pkg/controller/common/rollout/workloads/statefulset_scale_controller.go

@@ -126,7 +126,7 @@ func (s *StatefulSetScaleController) Initialize(ctx context.Context) (bool, erro
 		return false, nil
 	}
 
-	claimedBefore, err := s.claimStatefulSet(ctx, s.statefulSet, nil)
+	claimedBefore, err := s.claimStatefulSet(ctx, s.statefulSet)
 	if err != nil {
 		// nolint:nilerr
 		return false, nil

pkg/controller/common/rollout/workloads/statefulset_scale_integration_test.go

@@ -46,7 +46,7 @@ var _ = Describe("StatefulSet controller", func() {
 	BeforeEach(func() {
 		namespace = "rollout-ns"
 		name = "rollout1"
-		appRollout := v1beta1.AppRollout{ObjectMeta: metav1.ObjectMeta{Name: name}}
+		appRollout := v1beta1.AppRollout{TypeMeta: metav1.TypeMeta{APIVersion: v1beta1.SchemeGroupVersion.String(), Kind: v1beta1.AppRolloutKind}, ObjectMeta: metav1.ObjectMeta{Name: name}}
 		namespacedName = client.ObjectKey{Name: name, Namespace: namespace}
 
 		s = StatefulSetScaleController{

pkg/controller/core.oam.dev/v1alpha2/application/assemble/options.go

@@ -155,15 +155,21 @@ func PrepareWorkloadForRollout(rolloutComp string) WorkloadOption {
 					"kind", assembledWorkload.GetKind(), "instance name", assembledWorkload.GetName())
 				return nil
 			}
-		} else if assembledWorkload.GroupVersionKind().Group == appsv1.GroupName &&
-			assembledWorkload.GetKind() == reflect.TypeOf(appsv1.Deployment{}).Name() {
-			err := pv.SetBool(deploymentDisablePath, true)
-			if err != nil {
-				return err
+		}
+
+		if assembledWorkload.GroupVersionKind().Group == appsv1.GroupName {
+			switch assembledWorkload.GetKind() {
+			case reflect.TypeOf(appsv1.Deployment{}).Name():
+				if err := pv.SetBool(deploymentDisablePath, true); err != nil {
+					return err
+				}
+				klog.InfoS("we render a deployment assembledWorkload.paused on the first time",
+					"kind", assembledWorkload.GetKind(), "instance name", assembledWorkload.GetName())
+				return nil
+			case reflect.TypeOf(appsv1.StatefulSet{}).Name():
+				// TODO: Pause StatefulSet here.
+				return nil
 			}
-			klog.InfoS("we render a deployment assembledWorkload.paused on the first time",
-				"kind", assembledWorkload.GetKind(), "instance name", assembledWorkload.GetName())
-			return nil
 		}
 
 		klog.InfoS("we encountered an unknown resource, we don't know how to prepare it",

pkg/controller/core.oam.dev/v1alpha2/applicationrollout/helper.go

@@ -49,17 +49,26 @@ func RolloutWorkloadName(rolloutComp string) assemble.WorkloadOption {
 
 		// we hard code the behavior depends on the workload group/kind for now. The only in-place upgradable resources
 		// we support is cloneset/statefulset for now. We can easily add more later.
+		supportInplaceUpgrade := false
 		if w.GroupVersionKind().Group == v1alpha1.GroupVersion.Group {
-			if w.GetKind() == reflect.TypeOf(v1alpha1.CloneSet{}).Name() ||
-				w.GetKind() == reflect.TypeOf(v1alpha1.StatefulSet{}).Name() {
-				// we use the component name alone for those resources that do support in-place upgrade
-				klog.InfoS("we reuse the component name for resources that support in-place upgrade",
-					"GVK", w.GroupVersionKind(), "instance name", w.GetName())
-				// assemble use component name as workload name by default
-				// so no need to re-set name
-				return nil
+			if w.GetKind() == reflect.TypeOf(v1alpha1.CloneSet{}).Name() {
+				supportInplaceUpgrade = true
+			}
+		} else if w.GroupVersionKind().Group == appsv1.GroupName {
+			if w.GetKind() == reflect.TypeOf(appsv1.StatefulSet{}).Name() {
+				supportInplaceUpgrade = true
 			}
 		}
+
+		if supportInplaceUpgrade {
+			// we use the component name alone for those resources that do support in-place upgrade
+			klog.InfoS("we reuse the component name for resources that support in-place upgrade",
+				"GVK", w.GroupVersionKind(), "instance name", w.GetName())
+			// assemble use component name as workload name by default
+			// so no need to re-set name
+			return nil
+		}
+
 		// we assume that the rest of the resources do not support in-place upgrade
 		compRevName := w.GetLabels()[oam.LabelAppComponentRevision]
 		w.SetName(compRevName)
@@ -82,7 +91,7 @@ func HandleReplicas(ctx context.Context, rolloutComp string, c client.Client) as
 		// we hard code here, but we can easily support more types of workload by add more cases logic in switch
 		var replicasFieldPath string
 		switch u.GetKind() {
-		case reflect.TypeOf(v1alpha1.CloneSet{}).Name(), reflect.TypeOf(appsv1.Deployment{}).Name():
+		case reflect.TypeOf(v1alpha1.CloneSet{}).Name(), reflect.TypeOf(appsv1.Deployment{}).Name(), reflect.TypeOf(appsv1.StatefulSet{}).Name():
 			replicasFieldPath = "spec.replicas"
 		default:
 			klog.Errorf("rollout meet a workload we cannot support yet", "Kind", u.GetKind(), "name", u.GetName())