Support AdvancedStatefulSet in admission webhook

go.mod

@@ -47,7 +47,7 @@ require (
 	github.com/openshift/generic-admission-server v1.14.0
 	github.com/opentracing/opentracing-go v1.1.0 // indirect
 	github.com/pierrec/lz4 v2.0.5+incompatible // indirect
-	github.com/pingcap/advanced-statefulset v0.3.1
+	github.com/pingcap/advanced-statefulset v0.3.2
 	github.com/pingcap/check v0.0.0-20190102082844-67f458068fc8 // indirect
 	github.com/pingcap/errors v0.11.0
 	github.com/pingcap/kvproto v0.0.0-20191217072959-393e6c0fd4b7

go.sum

@@ -644,8 +644,8 @@ github.com/peterbourgon/diskv v2.0.1+incompatible h1:UBdAOUP5p4RWqPBg048CAvpKN+v
 github.com/peterbourgon/diskv v2.0.1+incompatible/go.mod h1:uqqh8zWWbv1HBMNONnaR/tNboyR3/BZd58JJSHlUSCU=
 github.com/pierrec/lz4 v2.0.5+incompatible h1:2xWsjqPFWcplujydGg4WmhC/6fZqK42wMM8aXeqhl0I=
 github.com/pierrec/lz4 v2.0.5+incompatible/go.mod h1:pdkljMzZIN41W+lC3N2tnIh5sFi+IEE17M5jbnwPHcY=
-github.com/pingcap/advanced-statefulset v0.3.1 h1:LxfAdpY2MV/b0MUlASYWjcPfUR161Xly1rA7oaIi684=
-github.com/pingcap/advanced-statefulset v0.3.1/go.mod h1:rg2p1v6AGsKhvEZi6Sm0YNYJCmdXdZZhQ6Sviei7Ivs=
+github.com/pingcap/advanced-statefulset v0.3.2 h1:cdnmWNaldoAyAWL/614Nr3hydnAzJEhSDMdIB6votZU=
+github.com/pingcap/advanced-statefulset v0.3.2/go.mod h1:rg2p1v6AGsKhvEZi6Sm0YNYJCmdXdZZhQ6Sviei7Ivs=
 github.com/pingcap/check v0.0.0-20190102082844-67f458068fc8 h1:USx2/E1bX46VG32FIw034Au6seQ2fY9NEILmNh/UlQg=
 github.com/pingcap/check v0.0.0-20190102082844-67f458068fc8/go.mod h1:B1+S9LNcuMyLH/4HMTViQOJevkGiik3wW2AN9zb2fNQ=
 github.com/pingcap/errors v0.11.0 h1:DCJQB8jrHbQ1VVlMFIrbj2ApScNNotVmkSNplu2yUt4=

pkg/webhook/pod/pd_deleter.go

@@ -14,7 +14,9 @@
 package pod
 
 import (
+	"github.com/pingcap/advanced-statefulset/pkg/apis/apps/v1/helper"
 	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
+	"github.com/pingcap/tidb-operator/pkg/features"
 	"github.com/pingcap/tidb-operator/pkg/label"
 	pdutil "github.com/pingcap/tidb-operator/pkg/manager/member"
 	operatorUtils "github.com/pingcap/tidb-operator/pkg/util"
@@ -36,7 +38,7 @@ func (pc *PodAdmissionControl) admitDeletePdPods(payload *admitPayload) *admissi
 
 	// If the pd pod is deleted by restarter, it is necessary to check former pd restart status
 	if _, exist := payload.pod.Annotations[label.AnnPodDeferDeleting]; exist {
-		existed, err := checkFormerPodRestartStatus(pc.kubeCli, v1alpha1.PDMemberType, payload.tc, namespace, ordinal, *payload.ownerStatefulSet.Spec.Replicas)
+		existed, err := checkFormerPodRestartStatus(pc.kubeCli, v1alpha1.PDMemberType, payload, ordinal)
 		if err != nil {
 			return util.ARFail(err)
 		}
@@ -191,12 +193,22 @@ func (pc *PodAdmissionControl) transferPDLeader(payload *admitPayload) *admissio
 		return util.ARFail(err)
 	}
 	tcName := payload.tc.Name
-	lastOrdinal := payload.tc.Status.PD.StatefulSet.Replicas - 1
 	var targetName string
-	if ordinal == lastOrdinal {
-		targetName = pdutil.PdPodName(tcName, 0)
+
+	if features.DefaultFeatureGate.Enabled(features.AdvancedStatefulSet) {
+		lastOrdinal := helper.GetMaxPodOrdinal(*payload.ownerStatefulSet.Spec.Replicas, payload.ownerStatefulSet)
+		if ordinal == lastOrdinal {
+			targetName = pdutil.PdPodName(tcName, helper.GetMinPodOrdinal(*payload.ownerStatefulSet.Spec.Replicas, payload.ownerStatefulSet))
+		} else {
+			targetName = pdutil.PdPodName(tcName, lastOrdinal)
+		}
 	} else {
-		targetName = pdutil.PdPodName(tcName, lastOrdinal)
+		lastOrdinal := payload.tc.Status.PD.StatefulSet.Replicas - 1
+		if ordinal == lastOrdinal {
+			targetName = pdutil.PdPodName(tcName, 0)
+		} else {
+			targetName = pdutil.PdPodName(tcName, lastOrdinal)
+		}
 	}
 
 	err = payload.pdClient.TransferPDLeader(targetName)

pkg/webhook/pod/pods.go

@@ -16,6 +16,7 @@ package pod
 import (
 	"encoding/json"
 	"fmt"
+	"github.com/pingcap/tidb-operator/pkg/features"
 	"time"
 
 	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
@@ -28,7 +29,6 @@ import (
 	"github.com/pingcap/tidb-operator/pkg/label"
 	memberUtils "github.com/pingcap/tidb-operator/pkg/manager/member"
 	"github.com/pingcap/tidb-operator/pkg/pdapi"
-	operatorUtils "github.com/pingcap/tidb-operator/pkg/util"
 	"github.com/pingcap/tidb-operator/pkg/webhook/util"
 	admission "k8s.io/api/admission/v1beta1"
 	"k8s.io/apimachinery/pkg/api/errors"
@@ -49,7 +49,8 @@ type PodAdmissionControl struct {
 }
 
 const (
-	stsControllerServiceAccounts = "system:serviceaccount:kube-system:statefulset-controller"
+	stsControllerServiceAccounts  = "system:serviceaccount:kube-system:statefulset-controller"
+	astsControllerServiceAccounts = "system:serviceaccount:pingcap:advanced-statefulset-controller"
 )
 
 func NewPodAdmissionControl(kubeCli kubernetes.Interface, operatorCli versioned.Interface, PdControl pdapi.PDControlInterface, extraServiceAccounts []string, evictRegionLeaderTimeout time.Duration) *PodAdmissionControl {
@@ -58,6 +59,9 @@ func NewPodAdmissionControl(kubeCli kubernetes.Interface, operatorCli versioned.
 	for _, sa := range extraServiceAccounts {
 		serviceAccounts.Insert(sa)
 	}
+	if features.DefaultFeatureGate.Enabled(features.AdvancedStatefulSet) {
+		serviceAccounts.Insert(astsControllerServiceAccounts)
+	}
 	EvictLeaderTimeout = evictRegionLeaderTimeout
 	return &PodAdmissionControl{
 		kubeCli:         kubeCli,
@@ -165,14 +169,10 @@ func (pc *PodAdmissionControl) admitDeletePods(name, namespace string) *admissio
 		return util.ARSuccess()
 	}
 
-	ordinal, err := operatorUtils.GetOrdinalFromPodName(name)
-	if err != nil {
-		return util.ARFail(err)
-	}
-
-	// If there was only one replica for this statefulset,admit to delete it.
-	if *ownerStatefulSet.Spec.Replicas == 1 && ordinal == 0 {
-		klog.Infof("tc[%s/%s]'s pd only have one pod[%s/%s],admit to delete it.", namespace, tcName, namespace, name)
+	// When AdvancedStatefulSet is enabled, the ordinal of the last pod in the statefulset could be a non-zero number,
+	// so we let the deleting request of the last pod pass when spec.replicas <= 1 and status.replicas equals 1
+	if *ownerStatefulSet.Spec.Replicas <= 1 && ownerStatefulSet.Status.Replicas == 1 {
+		klog.Infof("tc[%s/%s]'s statefulset only have one pod[%s/%s],admit to delete it.", namespace, tcName, namespace, name)
 		return util.ARSuccess()
 	}
 

pkg/webhook/pod/tikv_deleter.go

@@ -58,7 +58,7 @@ func (pc *PodAdmissionControl) admitDeleteTiKVPods(payload *admitPayload) *admis
 
 	// If the tikv pod is deleted by restarter, it is necessary to check former tikv restart status
 	if _, exist := payload.pod.Annotations[label.AnnPodDeferDeleting]; exist {
-		existed, err := checkFormerPodRestartStatus(pc.kubeCli, v1alpha1.TiKVMemberType, payload.tc, namespace, ordinal, *payload.ownerStatefulSet.Spec.Replicas)
+		existed, err := checkFormerPodRestartStatus(pc.kubeCli, v1alpha1.TiKVMemberType, payload, ordinal)
 		if err != nil {
 			return util.ARFail(err)
 		}
@@ -91,7 +91,7 @@ func (pc *PodAdmissionControl) admitDeleteTiKVPods(payload *admitPayload) *admis
 		}
 	}
 
-	if storeInfo == nil || storeInfo.Store == nil {
+	if !existed || storeInfo == nil || storeInfo.Store == nil {
 		klog.Infof("tc[%s/%s]'s tikv pod[%s/%s] can't be found store", namespace, tcName, namespace, name)
 		return pc.admitDeleteUselessTiKVPod(payload)
 	}

pkg/webhook/pod/util.go

@@ -15,6 +15,8 @@ package pod
 
 import (
 	"fmt"
+	"github.com/pingcap/advanced-statefulset/pkg/apis/apps/v1/helper"
+	"github.com/pingcap/tidb-operator/pkg/features"
 	"k8s.io/client-go/kubernetes"
 	"time"
 
@@ -141,16 +143,45 @@ func getOwnerStatefulSetForTiDBComponent(pod *core.Pod, kubeCli kubernetes.Inter
 
 // checkFormerPodRestartStatus checks whether there are any former pod is going to be restarted
 // return true if existed
-func checkFormerPodRestartStatus(kubeCli kubernetes.Interface, memberType v1alpha1.MemberType, tc *v1alpha1.TidbCluster, namespace string, ordinal int32, replicas int32) (bool, error) {
-	for i := replicas - 1; i > ordinal; i-- {
-		podName := memberUtil.MemberPodName(tc.Name, i, memberType)
+func checkFormerPodRestartStatus(kubeCli kubernetes.Interface, memberType v1alpha1.MemberType, payload *admitPayload, ordinal int32) (bool, error) {
+	namespace := payload.tc.Namespace
+	tc := payload.tc
+	replicas := *payload.ownerStatefulSet.Spec.Replicas
+
+	f := func(name string, ordinal int32, memberType v1alpha1.MemberType) (bool, error) {
+		podName := memberUtil.MemberPodName(tc.Name, ordinal, memberType)
 		pod, err := kubeCli.CoreV1().Pods(namespace).Get(podName, meta.GetOptions{})
 		if err != nil {
 			return false, err
 		}
 		if _, existed := pod.Annotations[label.AnnPodDeferDeleting]; existed {
 			return true, nil
 		}
+		return false, nil
+	}
+
+	if features.DefaultFeatureGate.Enabled(features.AdvancedStatefulSet) {
+		for k := range helper.GetPodOrdinals(replicas, payload.ownerStatefulSet) {
+			if k > ordinal {
+				existed, err := f(tc.Name, k, memberType)
+				if err != nil {
+					return false, err
+				}
+				if existed {
+					return true, nil
+				}
+			}
+		}
+	} else {
+		for i := replicas - 1; i > ordinal; i-- {
+			existed, err := f(tc.Name, i, memberType)
+			if err != nil {
+				return false, err
+			}
+			if existed {
+				return true, nil
+			}
+		}
 	}
 	return false, nil
 }

tests/asts.go

@@ -0,0 +1,181 @@
+// Copyright 2020 PingCAP, Inc.
+//
+// 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,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tests
+
+import (
+	"encoding/json"
+	"fmt"
+	"time"
+
+	"github.com/onsi/ginkgo"
+	"github.com/pingcap/advanced-statefulset/pkg/apis/apps/v1/helper"
+	"github.com/pingcap/tidb-operator/pkg/client/clientset/versioned"
+	"github.com/pingcap/tidb-operator/pkg/controller"
+	"github.com/pingcap/tidb-operator/pkg/label"
+	utilstatefulset "github.com/pingcap/tidb-operator/tests/e2e/util/statefulset"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/apimachinery/pkg/util/wait"
+	clientset "k8s.io/client-go/kubernetes"
+	"k8s.io/klog"
+	"k8s.io/kubernetes/test/e2e/framework"
+	e2esset "k8s.io/kubernetes/test/e2e/framework/statefulset"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+type AstsCase struct {
+	name        string
+	component   string // tikv,pd,tidb
+	replicas    int32
+	deleteSlots sets.Int32
+}
+
+func GetAstsCases() []AstsCase {
+	return []AstsCase{
+		{
+			name:        "Scaling in tikv from 5 to 3 by deleting pods 1 and 3",
+			component:   "tikv",
+			replicas:    3,
+			deleteSlots: sets.NewInt32(1, 3),
+		},
+		{
+			name:        "Scaling out tikv from 3 to 4 by adding pod 3",
+			component:   "tikv",
+			replicas:    4,
+			deleteSlots: sets.NewInt32(1),
+		},
+		{
+			name:        "Scaling tikv by adding pod 1 and deleting pod 2",
+			component:   "tikv",
+			replicas:    4,
+			deleteSlots: sets.NewInt32(2),
+		},
+		{
+			name:        "Scaling in tidb from 3 to 2 by deleting pod 1",
+			component:   "tidb",
+			replicas:    2,
+			deleteSlots: sets.NewInt32(1),
+		},
+		{
+			name:        "Scaling in tidb from 2 to 0",
+			component:   "tidb",
+			replicas:    0,
+			deleteSlots: sets.NewInt32(),
+		},
+		{
+			name:        "Scaling out tidb from 0 to 2 by adding pods 0 and 2",
+			component:   "tidb",
+			replicas:    2,
+			deleteSlots: sets.NewInt32(1),
+		},
+		{
+			name:        "Scaling tidb from 2 to 4 by deleting pods 2 and adding pods 3, 4 and 5",
+			component:   "tidb",
+			replicas:    4,
+			deleteSlots: sets.NewInt32(1, 2),
+		},
+		{
+			name:        "Scaling out pd from 3 to 5 by adding pods 3, 4",
+			component:   "pd",
+			replicas:    5,
+			deleteSlots: sets.NewInt32(),
+		},
+		{
+			name:        "Scaling in pd from 5 to 3 by deleting pods 0 and 3",
+			component:   "pd",
+			replicas:    3,
+			deleteSlots: sets.NewInt32(0, 3),
+		},
+		{
+			name:        "Scaling out pd from 3 to 5 by adding pods 5 and 6",
+			component:   "pd",
+			replicas:    5,
+			deleteSlots: sets.NewInt32(0, 3),
+		},
+	}
+}
+
+func AstsScalingTest(st AstsCase, clusterName, ns string, c, hc clientset.Interface, cli versioned.Interface, genericCli client.Client) {
+	ginkgo.By(st.name)
+	replicas := st.replicas
+	stsName := fmt.Sprintf("%s-%s", clusterName, st.component)
+
+	sts, err := hc.AppsV1().StatefulSets(ns).Get(stsName, metav1.GetOptions{})
+	framework.ExpectNoError(err)
+
+	oldPodList := e2esset.GetPodList(c, sts)
+
+	ginkgo.By(fmt.Sprintf("Scaling sts %s/%s to replicas %d and setting deleting pods to %v (old replicas: %d, old delete slots: %v)", ns, stsName, replicas, st.deleteSlots.List(), *sts.Spec.Replicas, helper.GetDeleteSlots(sts).List()))
+	tc, err := cli.PingcapV1alpha1().TidbClusters(ns).Get(clusterName, metav1.GetOptions{})
+	framework.ExpectNoError(err)
+	err = controller.GuaranteedUpdate(genericCli, tc, func() error {
+		if tc.Annotations == nil {
+			tc.Annotations = map[string]string{}
+		}
+		if st.component == "tikv" {
+			tc.Annotations[label.AnnTiKVDeleteSlots] = mustToString(st.deleteSlots)
+			tc.Spec.TiKV.Replicas = replicas
+		} else if st.component == "pd" {
+			tc.Annotations[label.AnnPDDeleteSlots] = mustToString(st.deleteSlots)
+			tc.Spec.PD.Replicas = replicas
+		} else if st.component == "tidb" {
+			tc.Annotations[label.AnnTiDBDeleteSlots] = mustToString(st.deleteSlots)
+			tc.Spec.TiDB.Replicas = replicas
+		} else {
+			return fmt.Errorf("unsupported component: %v", st.component)
+		}
+		return nil
+	})
+	framework.ExpectNoError(err)
+
+	ginkgo.By(fmt.Sprintf("Waiting for all pods of tidb cluster component %s (sts: %s/%s) are in desired state (replicas: %d, delete slots: %v)", st.component, ns, stsName, st.replicas, st.deleteSlots.List()))
+	err = wait.PollImmediate(time.Second*5, time.Minute*20, func() (bool, error) {
+		// check replicas and delete slots are synced
+		sts, err = hc.AppsV1().StatefulSets(ns).Get(stsName, metav1.GetOptions{})
+		if err != nil {
+			return false, nil
+		}
+		if *sts.Spec.Replicas != st.replicas {
+			klog.Infof("replicas of sts %s/%s is %d, expects %d", ns, stsName, sts.Spec.Replicas, st.replicas)
+			return false, nil
+		}
+		if !helper.GetDeleteSlots(sts).Equal(st.deleteSlots) {
+			klog.Infof("delete slots of sts %s/%s is %v, expects %v", ns, stsName, helper.GetDeleteSlots(sts).List(), st.deleteSlots.List())
+			return false, nil
+		}
+		// check all desired pods are running and ready
+		return utilstatefulset.IsAllDesiredPodsRunningAndReady(hc, sts), nil
+	})
+	framework.ExpectNoError(err)
+
+	ginkgo.By(fmt.Sprintf("Verify other pods of sts %s/%s should not be affected", ns, stsName))
+	newPodList := e2esset.GetPodList(c, sts)
+	framework.ExpectEqual(len(newPodList.Items), int(*sts.Spec.Replicas))
+	for _, newPod := range newPodList.Items {
+		for _, oldPod := range oldPodList.Items {
+			// if the pod is not new or deleted in scaling, it should not be affected
+			if oldPod.Name == newPod.Name && oldPod.UID != newPod.UID {
+				framework.Failf("pod %s/%s should not be affected (UID: %s, OLD UID: %s)", newPod.Namespace, newPod.Name, newPod.UID, oldPod.UID)
+			}
+		}
+	}
+}
+
+func mustToString(set sets.Int32) string {
+	b, err := json.Marshal(set.List())
+	if err != nil {
+		panic(err)
+	}
+	return string(b)
+}