support specified-delete in asts

Signed-off-by: Abner-1 <[email protected]>

pkg/controller/statefulset/stateful_set_control.go

@@ -45,6 +45,7 @@ import (
 	imagejobutilfunc "github.com/openkruise/kruise/pkg/util/imagejob/utilfunction"
 	"github.com/openkruise/kruise/pkg/util/inplaceupdate"
 	"github.com/openkruise/kruise/pkg/util/lifecycle"
+	"github.com/openkruise/kruise/pkg/util/specifieddelete"
 )
 
 // Realistic value for maximum in-flight requests when processing in parallel mode.
@@ -622,11 +623,18 @@ func (ssc *defaultStatefulSetControl) rollingUpdateStatefulsetPods(
 		}
 	}
 
+	specifiedDeletedPods := sets.NewString()
+	// handle specified deleted pod under maxUnavailable constrain
+	// NOTE: specified deletion is not constraint by partition setting
+	err = ssc.handleSpecifiedDeletedPods(set, status, currentRevision, updateRevision, replicas, maxUnavailable, unavailablePods, specifiedDeletedPods)
+	if err != nil {
+		return status, err
+	}
+
 	updateIndexes := sortPodsToUpdate(set.Spec.UpdateStrategy.RollingUpdate, updateRevision.Name, *set.Spec.Replicas, replicas)
 	klog.V(3).InfoS("Prepare to update pods indexes for StatefulSet", "statefulSet", klog.KObj(set), "podIndexes", updateIndexes)
 	// update pods in sequence
 	for _, target := range updateIndexes {
-
 		// the target is already up-to-date, go to next
 		if getPodRevision(replicas[target]) == updateRevision.Name {
 			continue
@@ -667,7 +675,7 @@ func (ssc *defaultStatefulSetControl) rollingUpdateStatefulsetPods(
 		}
 
 		// delete the Pod if it is not already terminating and does not match the update revision.
-		if !isTerminating(replicas[target]) {
+		if !specifiedDeletedPods.Has(replicas[target].Name) && !isTerminating(replicas[target]) {
 			// todo validate in-place for pub
 			inplacing, inplaceUpdateErr := ssc.inPlaceUpdatePod(set, replicas[target], updateRevision, revisions)
 			if inplaceUpdateErr != nil {
@@ -691,6 +699,48 @@ func (ssc *defaultStatefulSetControl) rollingUpdateStatefulsetPods(
 	return status, nil
 }
 
+func (ssc *defaultStatefulSetControl) handleSpecifiedDeletedPods(
+	set *appsv1beta1.StatefulSet,
+	status *appsv1beta1.StatefulSetStatus,
+	currentRevision *apps.ControllerRevision,
+	updateRevision *apps.ControllerRevision,
+	replicas []*v1.Pod,
+	maxUnavailable int,
+	unavailablePods,
+	specifiedDeletedPods sets.String) error {
+	for target := len(replicas) - 1; target >= 0; target-- {
+		if replicas[target] == nil {
+			continue
+		}
+		// the unavailable pods count exceed the maxUnavailable and the target is available, so we can't process it,
+		// why skip here rather than return?
+		// case: pod 0 ready, pod1 unready, pod 2 unready, pod3 ready, pod4 ready
+		// when maxUnavailable = 3, pod4 with specified deleted will be deleted but pod3 can't
+		// pod 2 and pod 1 can be deleted because they were unavailable
+		if len(unavailablePods) >= maxUnavailable && !unavailablePods.Has(replicas[target].Name) {
+			klog.V(4).InfoS("StatefulSet was waiting for unavailable Pods to update, blocked pod",
+				"statefulSet", klog.KObj(set), "unavailablePods", unavailablePods.List(), "blockedPod", klog.KObj(replicas[target]))
+			continue
+		}
+		if !specifieddelete.IsSpecifiedDelete(replicas[target]) {
+			continue
+		}
+		specifiedDeletedPods.Insert(replicas[target].Name)
+		if _, err := ssc.deletePod(set, replicas[target]); err != nil {
+			return err
+		}
+		// mark target as unavailable because it's updating
+		unavailablePods.Insert(replicas[target].Name)
+
+		if getPodRevision(replicas[target]) == currentRevision.Name {
+			status.CurrentReplicas--
+		} else if getPodRevision(replicas[target]) == updateRevision.Name {
+			status.UpdatedReplicas--
+		}
+	}
+	return nil
+}
+
 func (ssc *defaultStatefulSetControl) deletePod(set *appsv1beta1.StatefulSet, pod *v1.Pod) (bool, error) {
 	if set.Spec.Lifecycle != nil && lifecycle.IsPodHooked(set.Spec.Lifecycle.PreDelete, pod) {
 		markPodNotReady := set.Spec.Lifecycle.PreDelete.MarkPodNotReady

pkg/controller/statefulset/stateful_set_control_test.go

@@ -30,6 +30,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	apps "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/api/core/v1"
@@ -54,6 +55,7 @@ import (
 	utilpointer "k8s.io/utils/pointer"
 
 	appspub "github.com/openkruise/kruise/apis/apps/pub"
+	appsv1alpha1 "github.com/openkruise/kruise/apis/apps/v1alpha1"
 	appsv1beta1 "github.com/openkruise/kruise/apis/apps/v1beta1"
 	kruiseclientset "github.com/openkruise/kruise/pkg/client/clientset/versioned"
 	kruisefake "github.com/openkruise/kruise/pkg/client/clientset/versioned/fake"
@@ -2215,6 +2217,119 @@ func TestStatefulSetControlRollingUpdateBlockByMaxUnavailable(t *testing.T) {
 	}
 }
 
+func TestStatefulSetControlRollingUpdateWithSpecifiedDelete(t *testing.T) {
+	set := burst(newStatefulSet(6))
+	var partition int32 = 3
+	var maxUnavailable = intstr.FromInt(3)
+	set.Spec.UpdateStrategy = appsv1beta1.StatefulSetUpdateStrategy{
+		Type: apps.RollingUpdateStatefulSetStrategyType,
+		RollingUpdate: func() *appsv1beta1.RollingUpdateStatefulSetStrategy {
+			return &appsv1beta1.RollingUpdateStatefulSetStrategy{
+				Partition:       &partition,
+				MaxUnavailable:  &maxUnavailable,
+				PodUpdatePolicy: appsv1beta1.InPlaceIfPossiblePodUpdateStrategyType,
+			}
+		}(),
+	}
+
+	client := fake.NewSimpleClientset()
+	kruiseClient := kruisefake.NewSimpleClientset(set)
+	spc, _, ssc, stop := setupController(client, kruiseClient)
+	defer close(stop)
+	if err := scaleUpStatefulSetControl(set, ssc, spc, assertBurstInvariants); err != nil {
+		t.Fatal(err)
+	}
+	set, err := spc.setsLister.StatefulSets(set.Namespace).Get(set.Name)
+	if err != nil {
+		t.Fatal(err)
+	}
+	selector, err := metav1.LabelSelectorAsSelector(set.Spec.Selector)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// set pod 0 to specified delete
+	originalPods, err := spc.setPodSpecifiedDelete(set, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sort.Sort(ascendingOrdinal(originalPods))
+
+	// start to update
+	set.Spec.Template.Spec.Containers[0].Image = "foo"
+
+	// first update pod 5 only because pod 0 is specified deleted
+	if err = ssc.UpdateStatefulSet(context.TODO(), set, originalPods); err != nil {
+		t.Fatal(err)
+	}
+	pods, err := spc.podsLister.Pods(set.Namespace).List(selector)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// inplace update 5 and create 0
+	if err = ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
+		t.Fatal(err)
+	}
+	pods, err = spc.podsLister.Pods(set.Namespace).List(selector)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(pods) != 6 {
+		t.Fatalf("Expected create pods 5, got pods %v", pods)
+	}
+	sort.Sort(ascendingOrdinal(pods))
+	_, exist := pods[0].Labels[appsv1alpha1.SpecifiedDeleteKey]
+	assert.True(t, !exist)
+	// pod 0 is old image and pod 5/4 is new image
+	assert.Equal(t, pods[5].Spec.Containers[0].Image, "foo")
+	assert.Equal(t, pods[4].Spec.Containers[0].Image, "foo")
+	assert.Equal(t, pods[0].Spec.Containers[0].Image, "nginx")
+
+	// set pod 1/2/5 to specified deleted and pod 0/4/5 to ready
+	spc.setPodSpecifiedDelete(set, 1)
+	spc.setPodSpecifiedDelete(set, 2)
+	for i := 0; i < 6; i++ {
+		spc.setPodRunning(set, i)
+		spc.setPodReady(set, i)
+	}
+	originalPods, _ = spc.setPodSpecifiedDelete(set, 5)
+	sort.Sort(ascendingOrdinal(originalPods))
+
+	// create new pod for 1/2/5, do not update 3
+	if err = ssc.UpdateStatefulSet(context.TODO(), set, originalPods); err != nil {
+		t.Fatal(err)
+	}
+	pods, err = spc.podsLister.Pods(set.Namespace).List(selector)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// create new pods 5 and inplace update 3
+	if err = ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
+		t.Fatal(err)
+	}
+	pods, err = spc.podsLister.Pods(set.Namespace).List(selector)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sort.Sort(ascendingOrdinal(pods))
+	if len(pods) != 6 {
+		t.Fatalf("Expected create pods 5, got pods %v", pods)
+	}
+
+	_, exist = pods[5].Labels[appsv1alpha1.SpecifiedDeleteKey]
+	assert.True(t, !exist)
+	_, exist = pods[2].Labels[appsv1alpha1.SpecifiedDeleteKey]
+	assert.True(t, !exist)
+	_, exist = pods[1].Labels[appsv1alpha1.SpecifiedDeleteKey]
+	assert.True(t, !exist)
+	assert.Equal(t, pods[5].Spec.Containers[0].Image, "foo")
+	assert.Equal(t, pods[3].Spec.Containers[0].Image, "nginx")
+	assert.Equal(t, pods[2].Spec.Containers[0].Image, "nginx")
+	assert.Equal(t, pods[1].Spec.Containers[0].Image, "nginx")
+}
+
 func TestStatefulSetControlInPlaceUpdate(t *testing.T) {
 	set := burst(newStatefulSet(3))
 	var partition int32 = 1
@@ -3119,6 +3234,21 @@ func (om *fakeObjectManager) setPodTerminated(set *appsv1beta1.StatefulSet, ordi
 	return om.podsLister.Pods(set.Namespace).List(selector)
 }
 
+func (om *fakeObjectManager) setPodSpecifiedDelete(set *appsv1beta1.StatefulSet, ordinal int) ([]*v1.Pod, error) {
+	pod := newStatefulSetPod(set, ordinal)
+	if pod.Labels == nil {
+		pod.Labels = make(map[string]string)
+	}
+	pod.Labels[appsv1alpha1.SpecifiedDeleteKey] = "true"
+	fakeResourceVersion(pod)
+	om.podsIndexer.Update(pod)
+	selector, err := metav1.LabelSelectorAsSelector(set.Spec.Selector)
+	if err != nil {
+		return nil, err
+	}
+	return om.podsLister.Pods(set.Namespace).List(selector)
+}
+
 var _ StatefulPodControlObjectManager = &fakeObjectManager{}
 
 type fakeStatefulSetStatusUpdater struct {

test/e2e/apps/statefulset.go

@@ -1659,6 +1659,16 @@ var _ = SIGDescribe("StatefulSet", func() {
 				gomega.Expect(pods.Items[i].Labels["test-update"]).To(gomega.Equal("yes"))
 			}
 		})
+
+		/*
+			Testname: StatefulSet, Specified delete
+			Description: Specified delete pod MUST under maxUnavailable constrain.
+		*/
+		framework.ConformanceIt("should perform rolling updates and roll backs of template modifications", func() {
+			ginkgo.By("Creating a new StatefulSet")
+			ss := framework.NewStatefulSet("ss2", ns, headlessSvcName, 3, nil, nil, labels)
+			testWithSpecifiedDeleted(c, kc, ns, ss)
+		})
 	})
 
 	//ginkgo.Describe("Deploy clustered applications [Feature:StatefulSet] [Slow]", func() {
@@ -2799,3 +2809,84 @@ func waitForPVCCapacity(ctx context.Context, c clientset.Interface, kc kruisecli
 			return true, nil
 		})
 }
+
+// This function is used by two tests to test StatefulSet rollbacks: one using
+// PVCs and one using no storage.
+func testWithSpecifiedDeleted(c clientset.Interface, kc kruiseclientset.Interface, ns string, ss *appsv1beta1.StatefulSet) {
+	sst := framework.NewStatefulSetTester(c, kc)
+	*(ss.Spec.Replicas) = 4
+	ss, err := kc.AppsV1beta1().StatefulSets(ns).Create(context.TODO(), ss, metav1.CreateOptions{})
+	gomega.Expect(err).NotTo(gomega.HaveOccurred())
+	sst.WaitForRunningAndReady(*ss.Spec.Replicas, ss)
+	ss = sst.WaitForStatus(ss)
+	currentRevision, updateRevision := ss.Status.CurrentRevision, ss.Status.UpdateRevision
+	gomega.Expect(currentRevision).To(gomega.Equal(updateRevision),
+		fmt.Sprintf("StatefulSet %s/%s created with update revision %s not equal to current revision %s",
+			ss.Namespace, ss.Name, updateRevision, currentRevision))
+	pods := sst.GetPodList(ss)
+	for i := range pods.Items {
+		gomega.Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(gomega.Equal(currentRevision),
+			fmt.Sprintf("Pod %s/%s revision %s is not equal to current revision %s",
+				pods.Items[i].Namespace,
+				pods.Items[i].Name,
+				pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+				currentRevision))
+	}
+	sst.SortStatefulPods(pods)
+	specifiedDeletePod := func(idx int) {
+		oldUid := pods.Items[idx].UID
+		err = setPodSpecifiedDelete(c, ns, pods.Items[idx].Name)
+		gomega.Expect(err).NotTo(gomega.HaveOccurred())
+
+		ss = sst.WaitForStatus(ss)
+		pods = sst.GetPodList(ss)
+		sst.SortStatefulPods(pods)
+		gomega.Expect(oldUid).ShouldNot(gomega.Equal(pods.Items[idx].UID))
+	}
+	specifiedDeletePod(1)
+	newImage := NewNginxImage
+	oldImage := ss.Spec.Template.Spec.Containers[0].Image
+
+	ginkgo.By(fmt.Sprintf("Updating StatefulSet template: update image from %s to %s", oldImage, newImage))
+	gomega.Expect(oldImage).NotTo(gomega.Equal(newImage), "Incorrect test setup: should update to a different image")
+	var partition int32 = 2
+	ss, err = framework.UpdateStatefulSetWithRetries(kc, ns, ss.Name, func(update *appsv1beta1.StatefulSet) {
+		update.Spec.Template.Spec.Containers[0].Image = newImage
+		if update.Spec.UpdateStrategy.RollingUpdate == nil {
+			update.Spec.UpdateStrategy.RollingUpdate = &appsv1beta1.RollingUpdateStatefulSetStrategy{
+				Partition: &partition,
+			}
+		}
+	})
+	gomega.Expect(err).NotTo(gomega.HaveOccurred())
+	specifiedDeletePod(2)
+
+	ginkgo.By("Creating a new revision")
+	ss = sst.WaitForStatus(ss)
+	currentRevision, updateRevision = ss.Status.CurrentRevision, ss.Status.UpdateRevision
+	gomega.Expect(currentRevision).NotTo(gomega.Equal(updateRevision),
+		"Current revision should not equal update revision during rolling update")
+	specifiedDeletePod(1)
+	for i := range pods.Items {
+		if i >= int(partition) {
+			gomega.Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(gomega.Equal(updateRevision),
+				fmt.Sprintf("Pod %s/%s revision %s is not equal to updated revision %s",
+					pods.Items[i].Namespace,
+					pods.Items[i].Name,
+					pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+					updateRevision))
+		} else {
+			gomega.Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(gomega.Equal(currentRevision),
+				fmt.Sprintf("Pod %s/%s revision %s is not equal to current revision %s",
+					pods.Items[i].Namespace,
+					pods.Items[i].Name,
+					pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+					currentRevision))
+		}
+	}
+}
+
+func setPodSpecifiedDelete(c clientset.Interface, ns, name string) error {
+	_, err := c.CoreV1().Pods(ns).Patch(context.TODO(), name, types.StrategicMergePatchType, []byte(`{"metadata":{"labels":{"apps.kruise.io/specified-delete":"true"}}}`), metav1.PatchOptions{})
+	return err
+}