From e60428265c7b4a5cc462aa454a57edf5247217fb Mon Sep 17 00:00:00 2001
From: Song Gao <2695690803@qq.com>
Date: Wed, 4 Mar 2020 19:15:17 +0800
Subject: [PATCH] use tikv cli

---
 .../templates/scripts/_start_tikv.sh.tpl      |  5 +++
 .../admission/admission-webhook-rbac.yaml     |  2 +-
 pkg/manager/member/template.go                |  5 +++
 pkg/webhook/pod/pod_mutater.go                | 38 ++++++++++++++-----
 pkg/webhook/pod/util.go                       | 23 +++++++++++
 5 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/charts/tidb-cluster/templates/scripts/_start_tikv.sh.tpl b/charts/tidb-cluster/templates/scripts/_start_tikv.sh.tpl
index d4bb6c590f..1b88d47afb 100644
--- a/charts/tidb-cluster/templates/scripts/_start_tikv.sh.tpl
+++ b/charts/tidb-cluster/templates/scripts/_start_tikv.sh.tpl
@@ -39,6 +39,11 @@ ARGS="--pd={{ template "cluster.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}
diff --git a/charts/tidb-operator/templates/admission/admission-webhook-rbac.yaml b/charts/tidb-operator/templates/admission/admission-webhook-rbac.yaml
index 43d3082a8c..4d78b8e603 100644
--- a/charts/tidb-operator/templates/admission/admission-webhook-rbac.yaml
+++ b/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"]
diff --git a/pkg/manager/member/template.go b/pkg/manager/member/template.go
index f0e5a4c90a..e3ded6def4 100644
--- a/pkg/manager/member/template.go
+++ b/pkg/manager/member/template.go
@@ -225,6 +225,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}
diff --git a/pkg/webhook/pod/pod_mutater.go b/pkg/webhook/pod/pod_mutater.go
index 3f1d00058b..9f2d1598c9 100644
--- a/pkg/webhook/pod/pod_mutater.go
+++ b/pkg/webhook/pod/pod_mutater.go
@@ -16,7 +16,7 @@ package pod
 import (
 	"encoding/json"
 	"fmt"
-
+	"github.com/BurntSushi/toml"
 	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
 	"github.com/pingcap/tidb-operator/pkg/controller"
 	"github.com/pingcap/tidb-operator/pkg/features"
@@ -27,6 +27,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/klog"
 )
 
 func (pc *PodAdmissionControl) mutatePod(ar *admissionv1beta1.AdmissionRequest) *admissionv1beta1.AdmissionResponse {
@@ -35,7 +36,6 @@ func (pc *PodAdmissionControl) mutatePod(ar *admissionv1beta1.AdmissionRequest)
 		return util.ARFail(err)
 	}
 	original := pod.DeepCopy()
-
 	l := label.Label(pod.Labels)
 	if !l.IsManagedByTiDBOperator() {
 		return util.ARSuccess()
@@ -43,7 +43,6 @@ func (pc *PodAdmissionControl) mutatePod(ar *admissionv1beta1.AdmissionRequest)
 	if !l.IsTiKV() {
 		return util.ARSuccess()
 	}
-
 	tcName, exist := pod.Labels[label.InstanceLabelKey]
 	if !exist {
 		return util.ARSuccess()
@@ -83,12 +82,33 @@ func (pc *PodAdmissionControl) tikvHotRegionSchedule(tc *v1alpha1.TidbCluster, p
 		return nil
 	}
 
-	cmName := fmt.Sprintf("%s-autoscaling", controller.TiKVMemberName(tc.Name))
-	for _, v := range pod.Spec.Volumes {
-		if v.Name == "config" && v.ConfigMap != nil {
-			v.ConfigMap.LocalObjectReference = corev1.LocalObjectReference{
-				Name: cmName,
-			}
+	cmName := controller.TiKVMemberName(tc.Name)
+	cm, err := pc.kubeCli.CoreV1().ConfigMaps(tc.Namespace).Get(cmName, metav1.GetOptions{})
+	if err != nil {
+		klog.Infof("cm[%s/%s] found error,err %v", tc.Namespace, cmName, err)
+		return err
+	}
+	v, ok := cm.Data["config-file"]
+	if !ok {
+		return fmt.Errorf("tc[%s/%s]'s tikv config[config-file] is missing", tc.Namespace, tc.Name)
+	}
+	config := &v1alpha1.TiKVConfig{}
+	err = toml.Unmarshal([]byte(v), config)
+	if err != nil {
+		return err
+	}
+	if config.Server == nil {
+		config.Server = &v1alpha1.TiKVServerConfig{}
+	}
+	if config.Server.Labels == nil {
+		config.Server.Labels = map[string]string{}
+	}
+	// TODO: add document to explain the hot region label
+	config.Server.Labels["specialUse"] = "hotRegion"
+	for id, c := range pod.Spec.Containers {
+		if c.Name == "tikv" {
+			appendExtraLabelsENVForTiKV(config.Server.Labels, &c)
+			pod.Spec.Containers[id] = c
 			break
 		}
 	}
diff --git a/pkg/webhook/pod/util.go b/pkg/webhook/pod/util.go
index d6252c8f2e..16ea0cdaf8 100644
--- a/pkg/webhook/pod/util.go
+++ b/pkg/webhook/pod/util.go
@@ -180,3 +180,26 @@ func checkFormerPodRestartStatus(kubeCli kubernetes.Interface, memberType v1alph
 	}
 	return false, nil
 }
+
+func appendExtraLabelsENVForTiKV(labels map[string]string, container *core.Container) {
+	s := ""
+	for k, v := range labels {
+		s = fmt.Sprintf("%s,%s", s, fmt.Sprintf("%s=%s", k, v))
+	}
+	s = s[1:]
+	existed := false
+	for id, env := range container.Env {
+		if env.Name == "STORE_LABELS" {
+			env.Value = fmt.Sprintf("%s,%s", env.Value, s)
+			container.Env[id] = env
+			existed = true
+			break
+		}
+	}
+	if !existed {
+		container.Env = append(container.Env, core.EnvVar{
+			Name:  "STORE_LABELS",
+			Value: s,
+		})
+	}
+}