use tikv cli

pingcap · Mar 4, 2020 · e604282 · e604282
1 parent 8bffec2
commit e604282
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,15 +36,13 @@ 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()
 	}
 	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
@@ -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,
+		})
+	}
+}