Manually cp 1833 and 1801 (#1872)

* Manage hot region label for the tikv created by auto-scaler (#1801)

* Manage hot region label for the tikv created by auto-scaler

Co-authored-by: DanielZhangQD <36026334+DanielZhangQD@users.noreply.github.com>

* Fix hot region label setting for tikv auto-scaling (#1833)

* mutate

* fix log

* add admission configuration

* remove useless log

* format by comment

* use tikv cli

* remove useless code

* remove cmlister

* fix lint

* fix tpl

Co-authored-by: DanielZhangQD <36026334+DanielZhangQD@users.noreply.github.com>
Co-authored-by: Yecheng Fu <fuyecheng@pingcap.com>

charts/tidb-cluster/templates/scripts/_start_tikv.sh.tpl

@@ -39,6 +39,8 @@ ARGS="--pd={{ template "cluster.scheme" . }}://${CLUSTER_NAME}-pd:2379 \
 --config=/etc/tikv/tikv.toml
 "
 
+{{ .Values.tikv.postArgScript }}
+
 echo "starting tikv-server ..."
 echo "/tikv-server ${ARGS}"
 exec /tikv-server ${ARGS}

charts/tidb-cluster/values.yaml

@@ -284,6 +284,13 @@ tikv:
   # After waiting for 5 minutes, TiDB Operator creates a new TiKV node if this TiKV node is still down.
   # maxFailoverCount is used to configure the maximum number of TiKV nodes that TiDB Operator can create when failover occurs.
   maxFailoverCount: 3
+  # postArgscript is the script executed after the normal tikv instance start args is built,
+  # it is recommended to modify the args constructor logic if you have any special needs.
+  postArgScript: |
+    if [ ! -z "${STORE_LABELS:-}" ]; then
+      LABELS=" --labels ${STORE_LABELS} "
+      ARGS="${ARGS}${LABELS}"
+    fi
 
 tidb:
   # Please refer to https://github.com/pingcap/tidb/blob/master/config/config.toml.example for the default

charts/tidb-operator/templates/admission/admission-webhook-rbac.yaml

@@ -28,7 +28,7 @@ rules:
     resources: ["pods"]
     verbs: ["get", "list", "watch", "update"]
   - apiGroups: [""]
-    resources: ["secrets"]
+    resources: ["secrets","configmaps"]
     verbs: ["get", "list"]
   - apiGroups: [""]
     resources: ["events"]

charts/tidb-operator/templates/admission/admission-webhook-registration.yaml

@@ -162,4 +162,41 @@ webhooks:
         apiVersions: ["v1alpha1"]
         resources: ["tidbclusters"]
 {{- end }}
+---
+{{- if .Values.admissionWebhook.mutation.pods }}
+apiVersion: admissionregistration.k8s.io/v1beta1
+kind: MutatingWebhookConfiguration
+metadata:
+  name: mutation-tidb-pod-webhook-cfg
+  labels:
+    app.kubernetes.io/name: {{ template "chart.name" . }}
+    app.kubernetes.io/managed-by: {{ .Release.Service }}
+    app.kubernetes.io/instance: {{ .Release.Name }}
+    app.kubernetes.io/component: admission-webhook
+    helm.sh/chart: {{ .Chart.Name }}-{{ .Chart.Version | replace "+"  "_" }}
+webhooks:
+  - name: podadmission.tidb.pingcap.com
+    {{- if semverCompare ">=1.15-0" .Capabilities.KubeVersion.GitVersion }}
+    objectSelector:
+      matchLabels:
+        "app.kubernetes.io/managed-by": "tidb-operator"
+        "app.kubernetes.io/name": "tidb-cluster"
+    {{- end }}
+    failurePolicy: {{ .Values.admissionWebhook.failurePolicy.mutation | default "Ignore" }}
+    clientConfig:
+      service:
+        name: kubernetes
+        namespace: default
+        path: "/apis/admission.tidb.pingcap.com/v1alpha1/mutatingreviews"
+      {{- if .Values.admissionWebhook.cabundle }}
+      caBundle: {{ .Values.admissionWebhook.cabundle | b64enc }}
+      {{- else }}
+      caBundle: null
+      {{- end }}
+    rules:
+      - operations: ["CREATE"]
+        apiGroups: [""]
+        apiVersions: ["v1"]
+        resources: ["pods"]
+{{- end }}
 {{- end }}

charts/tidb-operator/values.yaml

@@ -197,6 +197,9 @@ admissionWebhook:
     pingcapResources: false
   ## mutation webhook would mutate the given request for the specific resource and operation
   mutation:
+    ## pods mutation hook would mutate the pod. Currently It is used for TiKV Auto-Scaling.
+    ## refer to https://github.com/pingcap/tidb-operator/issues/1651
+    pods: true
     ## defaulting hook set default values for the the resources under pingcap.com group
     pingcapResources: true
   ## failurePolicy are applied to ValidatingWebhookConfiguration which affect tidb-admission-webhook

pkg/autoscaler/autoscaler/tidb_autoscaler.go

@@ -46,13 +46,13 @@ func (am *autoScalerManager) syncTiDB(tc *v1alpha1.TidbCluster, tac *v1alpha1.Ti
 	if targetReplicas == tc.Spec.TiDB.Replicas {
 		return nil
 	}
-	return syncTiDBAfterCalculated(tc, tac, currentReplicas, targetReplicas)
+	return syncTiDBAfterCalculated(tc, tac, currentReplicas, targetReplicas, sts)
 }
 
 // syncTiDBAfterCalculated would check the Consecutive count to avoid jitter, and it would also check the interval
 // duration between each auto-scaling. If either of them is not meet, the auto-scaling would be rejected.
 // If the auto-scaling is permitted, the timestamp would be recorded and the Consecutive count would be zeroed.
-func syncTiDBAfterCalculated(tc *v1alpha1.TidbCluster, tac *v1alpha1.TidbClusterAutoScaler, currentReplicas, recommendedReplicas int32) error {
+func syncTiDBAfterCalculated(tc *v1alpha1.TidbCluster, tac *v1alpha1.TidbClusterAutoScaler, currentReplicas, recommendedReplicas int32, sts *appsv1.StatefulSet) error {
 	intervalSeconds := tac.Spec.TiDB.ScaleInIntervalSeconds
 	if recommendedReplicas > currentReplicas {
 		intervalSeconds = tac.Spec.TiDB.ScaleOutIntervalSeconds
@@ -64,13 +64,14 @@ func syncTiDBAfterCalculated(tc *v1alpha1.TidbCluster, tac *v1alpha1.TidbCluster
 	if !ableToScale {
 		return nil
 	}
-	updateTcTiDBAnnIfScale(tac)
-	tc.Spec.TiDB.Replicas = recommendedReplicas
-	return nil
+	return updateTcTiDBIfScale(tc, tac, recommendedReplicas)
 }
 
-func updateTcTiDBAnnIfScale(tac *v1alpha1.TidbClusterAutoScaler) {
+// Currently we didnt' record the auto-scaling out slot for tidb, because it is pointless for now.
+func updateTcTiDBIfScale(tc *v1alpha1.TidbCluster, tac *v1alpha1.TidbClusterAutoScaler, recommendedReplicas int32) error {
 	tac.Annotations[label.AnnTiDBLastAutoScalingTimestamp] = fmt.Sprintf("%d", time.Now().Unix())
+	tc.Spec.TiDB.Replicas = recommendedReplicas
+	return nil
 }
 
 func calculateTidbMetrics(tac *v1alpha1.TidbClusterAutoScaler, sts *appsv1.StatefulSet, instances []string) (int32, error) {

pkg/autoscaler/autoscaler/tikv_autoscaler.go

@@ -17,6 +17,7 @@ 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/autoscaler/autoscaler/calculate"
 	"github.com/pingcap/tidb-operator/pkg/label"
@@ -46,15 +47,15 @@ func (am *autoScalerManager) syncTiKV(tc *v1alpha1.TidbCluster, tac *v1alpha1.Ti
 	if targetReplicas == tc.Spec.TiKV.Replicas {
 		return nil
 	}
-	return syncTiKVAfterCalculated(tc, tac, currentReplicas, targetReplicas)
+	return syncTiKVAfterCalculated(tc, tac, currentReplicas, targetReplicas, sts)
 }
 
 // syncTiKVAfterCalculated would check the Consecutive count to avoid jitter, and it would also check the interval
 // duration between each auto-scaling. If either of them is not meet, the auto-scaling would be rejected.
 // If the auto-scaling is permitted, the timestamp would be recorded and the Consecutive count would be zeroed.
 // The currentReplicas of TiKV calculated in auto-scaling is the count of the StateUp TiKV instance, so we need to
 // add the number of other state tikv instance replicas when we update the TidbCluster.Spec.TiKV.Replicas
-func syncTiKVAfterCalculated(tc *v1alpha1.TidbCluster, tac *v1alpha1.TidbClusterAutoScaler, currentReplicas, recommendedReplicas int32) error {
+func syncTiKVAfterCalculated(tc *v1alpha1.TidbCluster, tac *v1alpha1.TidbClusterAutoScaler, currentReplicas, recommendedReplicas int32, sts *appsv1.StatefulSet) error {
 
 	intervalSeconds := tac.Spec.TiKV.ScaleInIntervalSeconds
 	if recommendedReplicas > tc.Spec.TiKV.Replicas {
@@ -67,9 +68,7 @@ func syncTiKVAfterCalculated(tc *v1alpha1.TidbCluster, tac *v1alpha1.TidbCluster
 	if !ableToScale {
 		return nil
 	}
-	updateTcTiKVAnnIfScale(tac)
-	tc.Spec.TiKV.Replicas = recommendedReplicas
-	return nil
+	return updateTcTiKVIfScale(tc, tac, currentReplicas, recommendedReplicas, sts)
 }
 
 //TODO: fetch tikv instances info from pdapi in future
@@ -83,8 +82,25 @@ func filterTiKVInstances(tc *v1alpha1.TidbCluster) []string {
 	return instances
 }
 
-func updateTcTiKVAnnIfScale(tac *v1alpha1.TidbClusterAutoScaler) {
+// we record the auto-scaling out slot for tikv, in order to add special hot labels when they are created
+func updateTcTiKVIfScale(tc *v1alpha1.TidbCluster, tac *v1alpha1.TidbClusterAutoScaler, currentReplicas, recommendedReplicas int32, sts *appsv1.StatefulSet) error {
 	tac.Annotations[label.AnnTiKVLastAutoScalingTimestamp] = fmt.Sprintf("%d", time.Now().Unix())
+	if recommendedReplicas > currentReplicas {
+		newlyScaleOutOrdinalSets := helper.GetPodOrdinals(recommendedReplicas, sts).Difference(helper.GetPodOrdinals(currentReplicas, sts))
+		if newlyScaleOutOrdinalSets.Len() > 0 {
+			if tc.Annotations == nil {
+				tc.Annotations = map[string]string{}
+			}
+			existed := operatorUtils.GetAutoScalingOutSlots(tc, v1alpha1.TiKVMemberType)
+			v, err := operatorUtils.Encode(newlyScaleOutOrdinalSets.Union(existed).List())
+			if err != nil {
+				return err
+			}
+			tc.Annotations[label.AnnTiKVAutoScalingOutOrdinals] = v
+		}
+	}
+	tc.Spec.TiKV.Replicas = recommendedReplicas
+	return nil
 }
 
 func calculateTikvMetrics(tac *v1alpha1.TidbClusterAutoScaler, sts *appsv1.StatefulSet, instances []string) (int32, error) {

pkg/label/label.go

@@ -114,6 +114,10 @@ const (
 	AnnAutoScalingTargetName = "auto-scaling.tidb.pingcap.com/target-name"
 	// AnnAutoScalingTargetNamespace describes the target TidbCluster Ref Namespace for the TidbCluserAutoScaler
 	AnnAutoScalingTargetNamespace = "auto-scaling.tidb.pingcap.com/target-namespace"
+	// AnnTiKVAutoScalingOutOrdinals describe the tikv pods' ordinal list which is created by auto-scaling out
+	AnnTiKVAutoScalingOutOrdinals = "tikv.tidb.pingcap.com/scale-out-ordinals"
+	// AnnTiDBAutoScalingOutOrdinals describe the tidb pods' ordinal list which is created by auto-scaling out
+	AnnTiDBAutoScalingOutOrdinals = "tidb.tidb.pingcap.com/scale-out-ordinals"
 
 	// PDLabelVal is PD label value
 	PDLabelVal string = "pd"

pkg/manager/member/pd_scaler.go

@@ -47,7 +47,7 @@ func (psd *pdScaler) Scale(tc *v1alpha1.TidbCluster, oldSet *apps.StatefulSet, n
 	} else if scaling < 0 {
 		return psd.ScaleIn(tc, oldSet, newSet)
 	}
-	return nil
+	return psd.SyncAutoScalerAnn(tc, oldSet)
 }
 
 func (psd *pdScaler) ScaleOut(tc *v1alpha1.TidbCluster, oldSet *apps.StatefulSet, newSet *apps.StatefulSet) error {
@@ -167,6 +167,10 @@ func (psd *pdScaler) ScaleIn(tc *v1alpha1.TidbCluster, oldSet *apps.StatefulSet,
 	return nil
 }
 
+func (psd *pdScaler) SyncAutoScalerAnn(tc *v1alpha1.TidbCluster, actual *apps.StatefulSet) error {
+	return nil
+}
+
 type fakePDScaler struct{}
 
 // NewFakePDScaler returns a fake Scaler
@@ -192,3 +196,7 @@ func (fsd *fakePDScaler) ScaleIn(_ *v1alpha1.TidbCluster, oldSet *apps.StatefulS
 	setReplicasAndDeleteSlots(newSet, *oldSet.Spec.Replicas-1, nil)
 	return nil
 }
+
+func (fsd *fakePDScaler) SyncAutoScalerAnn(tc *v1alpha1.TidbCluster, actual *apps.StatefulSet) error {
+	return nil
+}

pkg/manager/member/scaler.go

@@ -43,6 +43,8 @@ type Scaler interface {
 	ScaleOut(tc *v1alpha1.TidbCluster, actual *apps.StatefulSet, desired *apps.StatefulSet) error
 	// ScaleIn scales in the cluster
 	ScaleIn(tc *v1alpha1.TidbCluster, actual *apps.StatefulSet, desired *apps.StatefulSet) error
+	// SyncAutoScalerAnn would sync Ann created by AutoScaler
+	SyncAutoScalerAnn(tc *v1alpha1.TidbCluster, actual *apps.StatefulSet) error
 }
 
 type generalScaler struct {

pkg/manager/member/template.go

@@ -231,6 +231,11 @@ ARGS="--pd={{ .Scheme }}://${CLUSTER_NAME}-pd:2379 \
 --config=/etc/tikv/tikv.toml
 "
 
+if [ ! -z "${STORE_LABELS:-}" ]; then
+  LABELS=" --labels ${STORE_LABELS} "
+  ARGS="${ARGS}${LABELS}"
+fi
+
 echo "starting tikv-server ..."
 echo "/tikv-server ${ARGS}"
 exec /tikv-server ${ARGS}

pkg/manager/member/tikv_member_manager.go

@@ -61,7 +61,8 @@ type tikvMemberManager struct {
 }
 
 // NewTiKVMemberManager returns a *tikvMemberManager
-func NewTiKVMemberManager(pdControl pdapi.PDControlInterface,
+func NewTiKVMemberManager(
+	pdControl pdapi.PDControlInterface,
 	setControl controller.StatefulSetControlInterface,
 	svcControl controller.ServiceControlInterface,
 	certControl controller.CertControlInterface,

pkg/manager/member/tikv_scaler.go

@@ -18,10 +18,12 @@ import (
 	"strconv"
 	"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/controller"
 	"github.com/pingcap/tidb-operator/pkg/label"
 	"github.com/pingcap/tidb-operator/pkg/pdapi"
+	"github.com/pingcap/tidb-operator/pkg/util"
 	apps "k8s.io/api/apps/v1"
 	corelisters "k8s.io/client-go/listers/core/v1"
 	"k8s.io/klog"
@@ -48,7 +50,8 @@ func (tsd *tikvScaler) Scale(tc *v1alpha1.TidbCluster, oldSet *apps.StatefulSet,
 	} else if scaling < 0 {
 		return tsd.ScaleIn(tc, oldSet, newSet)
 	}
-	return nil
+	// we only sync auto scaler annotations when we are finishing syncing scaling
+	return tsd.SyncAutoScalerAnn(tc, oldSet)
 }
 
 func (tsd *tikvScaler) ScaleOut(tc *v1alpha1.TidbCluster, oldSet *apps.StatefulSet, newSet *apps.StatefulSet) error {
@@ -190,6 +193,32 @@ func (tsd *tikvScaler) ScaleIn(tc *v1alpha1.TidbCluster, oldSet *apps.StatefulSe
 	return fmt.Errorf("TiKV %s/%s not found in cluster", ns, podName)
 }
 
+// SyncAutoScalerAnn would reclaim the auto-scaling out slots if the target pod is no longer existed
+func (tsd *tikvScaler) SyncAutoScalerAnn(tc *v1alpha1.TidbCluster, actual *apps.StatefulSet) error {
+	currentScalingSlots := util.GetAutoScalingOutSlots(tc, v1alpha1.TiKVMemberType)
+	if currentScalingSlots.Len() < 1 {
+		return nil
+	}
+	currentOrdinals := helper.GetPodOrdinals(tc.Spec.TiKV.Replicas, actual)
+
+	// reclaim the auto-scaling out slots if the target pod is no longer existed
+	if !currentOrdinals.HasAll(currentScalingSlots.List()...) {
+		reclaimedSlots := currentScalingSlots.Difference(currentOrdinals)
+		currentScalingSlots = currentScalingSlots.Delete(reclaimedSlots.List()...)
+		if currentScalingSlots.Len() < 1 {
+			delete(tc.Annotations, label.AnnTiKVAutoScalingOutOrdinals)
+			return nil
+		}
+		v, err := util.Encode(currentScalingSlots.List())
+		if err != nil {
+			return err
+		}
+		tc.Annotations[label.AnnTiKVAutoScalingOutOrdinals] = v
+		return nil
+	}
+	return nil
+}
+
 type fakeTiKVScaler struct{}
 
 // NewFakeTiKVScaler returns a fake tikv Scaler
@@ -215,3 +244,7 @@ func (fsd *fakeTiKVScaler) ScaleIn(_ *v1alpha1.TidbCluster, oldSet *apps.Statefu
 	setReplicasAndDeleteSlots(newSet, *oldSet.Spec.Replicas-1, nil)
 	return nil
 }
+
+func (fsd *fakeTiKVScaler) SyncAutoScalerAnn(tc *v1alpha1.TidbCluster, actual *apps.StatefulSet) error {
+	return nil
+}

pkg/manager/member/utils.go

@@ -24,6 +24,7 @@ import (
 	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
 	"github.com/pingcap/tidb-operator/pkg/controller"
 	"github.com/pingcap/tidb-operator/pkg/label"
+	"github.com/pingcap/tidb-operator/pkg/util"
 	apps "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	apiequality "k8s.io/apimachinery/pkg/api/equality"
@@ -71,7 +72,7 @@ func statefulSetIsUpgrading(set *apps.StatefulSet) bool {
 
 // SetStatefulSetLastAppliedConfigAnnotation set last applied config to Statefulset's annotation
 func SetStatefulSetLastAppliedConfigAnnotation(set *apps.StatefulSet) error {
-	setApply, err := encode(set.Spec)
+	setApply, err := util.Encode(set.Spec)
 	if err != nil {
 		return err
 	}
@@ -97,14 +98,6 @@ func GetLastAppliedConfig(set *apps.StatefulSet) (*apps.StatefulSetSpec, *corev1
 	return spec, &spec.Template.Spec, nil
 }
 
-func encode(obj interface{}) (string, error) {
-	b, err := json.Marshal(obj)
-	if err != nil {
-		return "", err
-	}
-	return string(b), nil
-}
-
 // statefulSetEqual compares the new Statefulset's spec with old Statefulset's last applied config
 func statefulSetEqual(new apps.StatefulSet, old apps.StatefulSet) bool {
 	if !apiequality.Semantic.DeepEqual(new.Annotations, old.Annotations) {
@@ -216,6 +209,10 @@ func MarshalTOML(v interface{}) ([]byte, error) {
 	return data, nil
 }
 
+func UnmarshalTOML(b []byte, obj interface{}) error {
+	return toml.Unmarshal(b, obj)
+}
+
 func Sha256Sum(v interface{}) (string, error) {
 	data, err := json.Marshal(v)
 	if err != nil {