Skip to content

Commit

Permalink
check the StatefulSet in the incoming request and add UTs (pingcap#2997)
Browse files Browse the repository at this point in the history
Co-authored-by: DanielZhangQD <[email protected]>
  • Loading branch information
cofyc and DanielZhangQD authored Jul 22, 2020
1 parent 741b18e commit 80f74d7
Show file tree
Hide file tree
Showing 3 changed files with 319 additions and 22 deletions.
40 changes: 23 additions & 17 deletions pkg/webhook/statefulset/statefulset.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,10 @@
package statefulset

import (
"errors"
"fmt"
"k8s.io/klog"
"strconv"

asappsv1 "github.com/pingcap/advanced-statefulset/client/apis/apps/v1"
asapps "github.com/pingcap/advanced-statefulset/client/apis/apps/v1"
"github.com/pingcap/tidb-operator/pkg/client/clientset/versioned"
"github.com/pingcap/tidb-operator/pkg/controller"
"github.com/pingcap/tidb-operator/pkg/features"
Expand All @@ -29,16 +27,13 @@ import (
apps "k8s.io/api/apps/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/runtime"
"k8s.io/klog"
)

var (
deserializer runtime.Decoder
deserializer runtime.Decoder = util.Codecs.UniversalDeserializer()
)

func init() {
deserializer = util.GetCodec()
}

type StatefulSetAdmissionControl struct {
// operator client interface
operatorCli versioned.Interface
Expand All @@ -56,14 +51,14 @@ func (sc *StatefulSetAdmissionControl) AdmitStatefulSets(ar *admission.Admission
namespace := ar.Namespace
expectedGroup := "apps"
if features.DefaultFeatureGate.Enabled(features.AdvancedStatefulSet) {
expectedGroup = asappsv1.GroupName
expectedGroup = asapps.GroupName
}
apiVersion := ar.Resource.Version
setResource := metav1.GroupVersionResource{Group: expectedGroup, Version: apiVersion, Resource: "statefulsets"}

klog.Infof("admit %s [%s/%s]", setResource, namespace, name)

stsObjectMeta, stsPartition, err := getStsAttributes(ar.OldObject.Raw)
stsObjectMeta, stsPartition, err := getStsAttributes(ar.Object.Raw)
if err != nil {
err = fmt.Errorf("statefulset %s/%s, decode request failed, err: %v", namespace, name, err)
klog.Error(err)
Expand Down Expand Up @@ -93,11 +88,11 @@ func (sc *StatefulSetAdmissionControl) AdmitStatefulSets(ar *admission.Admission
return util.ARFail(err)
}

var partitionStr string
partitionStr = tc.Annotations[label.AnnTiDBPartition]
annKey := label.AnnTiDBPartition
if l.IsTiKV() {
partitionStr = tc.Annotations[label.AnnTiKVPartition]
annKey = label.AnnTiKVPartition
}
partitionStr := tc.Annotations[annKey]

if len(partitionStr) == 0 {
return util.ARSuccess()
Expand All @@ -111,17 +106,28 @@ func (sc *StatefulSetAdmissionControl) AdmitStatefulSets(ar *admission.Admission
}

if stsPartition != nil {
if *stsPartition > 0 && *stsPartition <= int32(partition) {
klog.Infof("statefulset %s/%s has been protect by partition %s annotations", namespace, name, partitionStr)
return util.ARFail(errors.New("protect by partition annotation"))
// only [partition from tc, INT32_MAX] are allowed
if *stsPartition < int32(partition) {
klog.Infof("statefulset %s/%s has been protected by partition annotation %q", namespace, name, partitionStr)
return util.ARFail(fmt.Errorf("protected by partition annotation (%s = %s) on the tidb cluster %s/%s", annKey, partitionStr, namespace, tcName))
}
klog.Infof("admit statefulset %s/%s update partition to %d, protect partition is %d", namespace, name, *stsPartition, partition)
}
return util.ARSuccess()
}

func getStsAttributes(data []byte) (*metav1.ObjectMeta, *int32, error) {
set := apps.StatefulSet{}
if !features.DefaultFeatureGate.Enabled(features.AdvancedStatefulSet) {
set := apps.StatefulSet{}
if _, _, err := deserializer.Decode(data, nil, &set); err != nil {
return nil, nil, err
}
if set.Spec.UpdateStrategy.RollingUpdate != nil {
return &(set.ObjectMeta), set.Spec.UpdateStrategy.RollingUpdate.Partition, nil
}
return &(set.ObjectMeta), nil, nil
}
set := asapps.StatefulSet{}
if _, _, err := deserializer.Decode(data, nil, &set); err != nil {
return nil, nil, err
}
Expand Down
285 changes: 285 additions & 0 deletions pkg/webhook/statefulset/statefulset_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,285 @@
// 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 statefulset

import (
"bytes"
"testing"

asapps "github.com/pingcap/advanced-statefulset/client/apis/apps/v1"
"github.com/pingcap/advanced-statefulset/client/apis/apps/v1/helper"
"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
"github.com/pingcap/tidb-operator/pkg/client/clientset/versioned/fake"
"github.com/pingcap/tidb-operator/pkg/features"
"github.com/pingcap/tidb-operator/pkg/webhook/util"
admission "k8s.io/api/admission/v1beta1"
apps "k8s.io/api/apps/v1"
v1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/runtime"
"k8s.io/utils/pointer"
)

type testcase struct {
name string
sts *apps.StatefulSet
tc *v1alpha1.TidbCluster
operation admission.Operation
wantAllowed bool
}

var (
ownerTCName = "foo"
validOwnerRefs = []metav1.OwnerReference{
{
APIVersion: "pingcap.com/v1alpha1",
Kind: "TidbCluster",
Name: ownerTCName,
Controller: pointer.BoolPtr(true),
},
}
tests = []testcase{
{
name: "non-update operation",
operation: admission.Delete,
wantAllowed: true,
},
{
name: "not tidb or tikv",
operation: admission.Update,
sts: &apps.StatefulSet{
TypeMeta: metav1.TypeMeta{
Kind: "StatefulSet",
APIVersion: "apps/v1",
},
ObjectMeta: metav1.ObjectMeta{},
},
wantAllowed: true,
},
{
name: "not owned by tidb cluster",
operation: admission.Update,
sts: &apps.StatefulSet{
TypeMeta: metav1.TypeMeta{
Kind: "StatefulSet",
APIVersion: "apps/v1",
},
ObjectMeta: metav1.ObjectMeta{
Labels: map[string]string{
"app.kubernetes.io/component": "tikv",
},
},
},
wantAllowed: true,
},
{
name: "owned by tidb cluster but the cluster does not exist",
operation: admission.Update,
sts: &apps.StatefulSet{
TypeMeta: metav1.TypeMeta{
Kind: "StatefulSet",
APIVersion: "apps/v1",
},
ObjectMeta: metav1.ObjectMeta{
Labels: map[string]string{
"app.kubernetes.io/component": "tikv",
},
OwnerReferences: validOwnerRefs,
},
},
wantAllowed: false,
},
{
name: "partition does not exist",
operation: admission.Update,
sts: &apps.StatefulSet{
TypeMeta: metav1.TypeMeta{
Kind: "StatefulSet",
APIVersion: "apps/v1",
},
ObjectMeta: metav1.ObjectMeta{
Labels: map[string]string{
"app.kubernetes.io/component": "tikv",
},
OwnerReferences: validOwnerRefs,
},
},
tc: &v1alpha1.TidbCluster{
ObjectMeta: metav1.ObjectMeta{
Name: ownerTCName,
Namespace: v1.NamespaceDefault,
Annotations: nil,
},
},
wantAllowed: true,
},
{
name: "partition is invald",
operation: admission.Update,
sts: &apps.StatefulSet{
TypeMeta: metav1.TypeMeta{
Kind: "StatefulSet",
APIVersion: "apps/v1",
},
ObjectMeta: metav1.ObjectMeta{
Labels: map[string]string{
"app.kubernetes.io/component": "tikv",
},
OwnerReferences: validOwnerRefs,
},
},
tc: &v1alpha1.TidbCluster{
ObjectMeta: metav1.ObjectMeta{
Name: ownerTCName,
Namespace: v1.NamespaceDefault,
Annotations: map[string]string{
"tidb.pingcap.com/tikv-partition": "invalid",
},
},
},
wantAllowed: false,
},
{
name: "partition allow",
operation: admission.Update,
sts: &apps.StatefulSet{
TypeMeta: metav1.TypeMeta{
Kind: "StatefulSet",
APIVersion: "apps/v1",
},
ObjectMeta: metav1.ObjectMeta{
Labels: map[string]string{
"app.kubernetes.io/component": "tikv",
},
OwnerReferences: validOwnerRefs,
},
Spec: apps.StatefulSetSpec{
UpdateStrategy: apps.StatefulSetUpdateStrategy{
Type: apps.RollingUpdateStatefulSetStrategyType,
RollingUpdate: &apps.RollingUpdateStatefulSetStrategy{
Partition: pointer.Int32Ptr(1),
},
},
},
},
tc: &v1alpha1.TidbCluster{
ObjectMeta: metav1.ObjectMeta{
Name: ownerTCName,
Namespace: v1.NamespaceDefault,
Annotations: map[string]string{
"tidb.pingcap.com/tikv-partition": "1",
},
},
},
wantAllowed: true,
},
{
name: "partition deny",
operation: admission.Update,
sts: &apps.StatefulSet{
TypeMeta: metav1.TypeMeta{
Kind: "StatefulSet",
APIVersion: "apps/v1",
},
ObjectMeta: metav1.ObjectMeta{
Labels: map[string]string{
"app.kubernetes.io/component": "tikv",
},
OwnerReferences: validOwnerRefs,
},
Spec: apps.StatefulSetSpec{
UpdateStrategy: apps.StatefulSetUpdateStrategy{
Type: apps.RollingUpdateStatefulSetStrategyType,
RollingUpdate: &apps.RollingUpdateStatefulSetStrategy{
Partition: pointer.Int32Ptr(0),
},
},
},
},
tc: &v1alpha1.TidbCluster{
ObjectMeta: metav1.ObjectMeta{
Name: ownerTCName,
Namespace: v1.NamespaceDefault,
Annotations: map[string]string{
"tidb.pingcap.com/tikv-partition": "1",
},
},
},
wantAllowed: false,
},
}
)

func runTest(t *testing.T, tt testcase, asts bool) {
jsonInfo, ok := runtime.SerializerInfoForMediaType(util.Codecs.SupportedMediaTypes(), runtime.ContentTypeJSON)
if !ok {
t.Fatalf("unable to locate encoder -- %q is not a supported media type", runtime.ContentTypeJSON)
}

cli := fake.NewSimpleClientset()
ac := NewStatefulSetAdmissionControl(cli)
ar := &admission.AdmissionRequest{
Name: "foo",
Namespace: v1.NamespaceDefault,
Operation: tt.operation,
}
if tt.sts != nil {
buf := bytes.Buffer{}
if asts {
encoder := util.Codecs.EncoderForVersion(jsonInfo.Serializer, asapps.SchemeGroupVersion)
sts, err := helper.FromBuiltinStatefulSet(tt.sts)
if err != nil {
t.Fatal(err)
}
if err := encoder.Encode(sts, &buf); err != nil {
t.Fatal(err)
}
} else {
encoder := util.Codecs.EncoderForVersion(jsonInfo.Serializer, apps.SchemeGroupVersion)
if err := encoder.Encode(tt.sts, &buf); err != nil {
t.Fatal(err)
}
}
ar.Object = runtime.RawExtension{
Raw: buf.Bytes(),
}
}
if tt.tc != nil {
cli.PingcapV1alpha1().TidbClusters(tt.tc.Namespace).Create(tt.tc)
}
resp := ac.AdmitStatefulSets(ar)
if resp.Allowed != tt.wantAllowed {
t.Errorf("want allowed %v, got %v", tt.wantAllowed, resp.Allowed)
}
}

func TestStatefulSetAdmissionControl(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
runTest(t, tt, false)
})
}
}

func TestStatefulSetAdmissionControl_ASTS(t *testing.T) {
saved := features.DefaultFeatureGate.String()
features.DefaultFeatureGate.Set("AdvancedStatefulSet=true")
defer features.DefaultFeatureGate.Set(saved) // reset features on exit

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
runTest(t, tt, true)
})
}
}
Loading

0 comments on commit 80f74d7

Please sign in to comment.