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

@@ -645,8 +645,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,6 +14,7 @@
 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/label"
 	pdutil "github.com/pingcap/tidb-operator/pkg/manager/member"
@@ -36,7 +37,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,10 +192,11 @@ 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
+
+	lastOrdinal := helper.GetMaxPodOrdinal(*payload.ownerStatefulSet.Spec.Replicas, payload.ownerStatefulSet)
 	if ordinal == lastOrdinal {
-		targetName = pdutil.PdPodName(tcName, 0)
+		targetName = pdutil.PdPodName(tcName, helper.GetMinPodOrdinal(*payload.ownerStatefulSet.Spec.Replicas, payload.ownerStatefulSet))
 	} else {
 		targetName = pdutil.PdPodName(tcName, lastOrdinal)
 	}

pkg/webhook/pod/pods.go

@@ -19,20 +19,18 @@ import (
 	"time"
 
 	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
-	apps "k8s.io/api/apps/v1"
-
-	core "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/util/sets"
-
 	"github.com/pingcap/tidb-operator/pkg/client/clientset/versioned"
+	"github.com/pingcap/tidb-operator/pkg/features"
 	"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"
+	apps "k8s.io/api/apps/v1"
+	core "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/klog"
 )
@@ -49,7 +47,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 +57,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 +167,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

@@ -17,7 +17,9 @@ import (
 	"fmt"
 	"time"
 
+	"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"
 	memberUtil "github.com/pingcap/tidb-operator/pkg/manager/member"
 	"github.com/pingcap/tidb-operator/pkg/pdapi"
@@ -136,16 +138,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/e2e/tidbcluster/serial.go

@@ -107,7 +107,7 @@ var _ = ginkgo.Describe("[tidb-operator][Serial]", func() {
 	})
 
 	// tidb-operator with AdvancedStatefulSet feature enabled
-	ginkgo.Context("[Feature: AdvancedStatefulSet]", func() {
+	ginkgo.Context("[Feature: AdvancedStatefulSet][Feature: Webhook]", func() {
 		var ocfg *tests.OperatorConfig
 		var oa tests.OperatorActions
 		var genericCli client.Client
@@ -123,9 +123,12 @@ var _ = ginkgo.Describe("[tidb-operator][Serial]", func() {
 					"StableScheduling=true",
 					"AdvancedStatefulSet=true",
 				},
-				LogLevel:        "4",
-				ImagePullPolicy: v1.PullIfNotPresent,
-				TestMode:        true,
+				LogLevel:          "4",
+				ImagePullPolicy:   v1.PullIfNotPresent,
+				TestMode:          true,
+				WebhookEnabled:    true,
+				PodWebhookEnabled: true,
+				StsWebhookEnabled: false,
 			}
 			oa = tests.NewOperatorActions(cli, c, asCli, aggrCli, apiExtCli, tests.DefaultPollInterval, ocfg, e2econfig.TestConfig, nil, fw, f)
 			ginkgo.By("Installing CRDs")