add storeLabels setting (#527)

charts/tidb-cluster/templates/config/_pd-config.tpl

@@ -82,7 +82,11 @@ max-replicas = {{ .Values.pd.maxReplicas }}
 # The placement priorities is implied by the order of label keys.
 # For example, ["zone", "rack"] means that we should place replicas to
 # different zones first, then to different racks if we don't have enough zones.
+{{- if .Values.tikv.storeLabels }}
+location-labels = {{ toJson .Values.tikv.storeLabels }}
+{{- else }}
 location-labels = ["region", "zone", "rack", "host"]
+{{- end }}
 
 [label-property]
 # Do not assign region leaders to stores that have these tags.

charts/tidb-cluster/templates/tidb-cluster.yaml

@@ -67,6 +67,16 @@ spec:
     annotations:
 {{ toYaml .Values.tikv.annotations | indent 6 }}
   {{- end }}
+  {{- if .Values.tikv.storeLabels }}
+    storeLabels:
+{{ toYaml .Values.tikv.storeLabels | indent 4 }}
+  {{- else }}
+    storeLabels:
+      - region
+      - zone
+      - rack
+      - host
+  {{- end }}
   tidb:
     replicas: {{ .Values.tidb.replicas }}
     image: {{ .Values.tidb.image }}

charts/tidb-cluster/values.yaml

@@ -192,6 +192,16 @@ tikv:
   #   effect: "NoSchedule"
   annotations: {}
 
+  ## storeLabels is used to define store label keys
+  ## The label keys specified the location of a store.
+  ## In order to use the location awareness feature of TiKV, users have to label their k8s nodes with the same labels.
+  ## Note: current can not support these labels contains "/"
+  ## The placement priorities is implied by the order of label keys.
+  ## For example, ["zone", "rack"] means that we should place replicas to
+  ## different zones first, then to different racks if we don't have enough zones.
+  ## default value is ["region", "zone", "rack", "host"]
+  ## storeLabels: ["region", "zone", "rack", "host"]
+
   # block-cache used to cache uncompressed blocks, big block-cache can speed up read.
   # in normal cases should tune to 30%-50% tikv.resources.limits.memory
   # defaultcfBlockCacheSize: "1GB"

pkg/apis/pingcap.com/v1alpha1/types.go

@@ -144,6 +144,7 @@ type TiKVSpec struct {
 	StorageClassName string              `json:"storageClassName,omitempty"`
 	Tolerations      []corev1.Toleration `json:"tolerations,omitempty"`
 	Annotations      map[string]string   `json:"annotations,omitempty"`
+	StoreLabels      []string            `json:"storeLabels,omitempty"`
 }
 
 // TiKVPromGatewaySpec runs as a sidecar with TiKVSpec

pkg/manager/member/tikv_member_manager.go

@@ -537,8 +537,8 @@ func (tkmm *tikvMemberManager) setStoreLabelsForTiKV(tc *v1alpha1.TidbCluster) (
 		}
 
 		nodeName := pod.Spec.NodeName
-		ls, err := tkmm.getNodeLabels(nodeName)
-		if err != nil {
+		ls, err := tkmm.getNodeLabels(nodeName, tc.Spec.TiKV.StoreLabels)
+		if err != nil || len(ls) == 0 {
 			glog.Warningf("node: [%s] has no node labels, skipping set store labels for Pod: [%s/%s]", nodeName, ns, podName)
 			continue
 		}
@@ -559,28 +559,28 @@ func (tkmm *tikvMemberManager) setStoreLabelsForTiKV(tc *v1alpha1.TidbCluster) (
 	return setCount, nil
 }
 
-func (tkmm *tikvMemberManager) getNodeLabels(nodeName string) (map[string]string, error) {
+func (tkmm *tikvMemberManager) getNodeLabels(nodeName string, storeLabels []string) (map[string]string, error) {
 	node, err := tkmm.nodeLister.Get(nodeName)
 	if err != nil {
 		return nil, err
 	}
-	if ls := node.GetLabels(); ls != nil {
-		labels := map[string]string{}
-		if region, found := ls["region"]; found {
-			labels["region"] = region
-		}
-		if zone, found := ls["zone"]; found {
-			labels["zone"] = zone
-		}
-		if rack, found := ls["rack"]; found {
-			labels["rack"] = rack
+	labels := map[string]string{}
+	ls := node.GetLabels()
+	for _, storeLabel := range storeLabels {
+		if value, found := ls[storeLabel]; found {
+			labels[storeLabel] = value
+			continue
 		}
-		if host, found := ls[apis.LabelHostname]; found {
-			labels["host"] = host
+
+		// TODO after pd supports storeLabel containing slash character, these codes should be deleted
+		if storeLabel == "host" {
+			if host, found := ls[apis.LabelHostname]; found {
+				labels[storeLabel] = host
+			}
 		}
-		return labels, nil
+
 	}
-	return nil, fmt.Errorf("labels not found")
+	return labels, nil
 }
 
 // storeLabelsEqualNodeLabels compares store labels with node labels

pkg/manager/member/tikv_member_manager_test.go

@@ -488,6 +488,7 @@ func TestTiKVMemberManagerSetStoreLabelsForTiKV(t *testing.T) {
 	}
 	testFn := func(test *testcase, t *testing.T) {
 		tc := newTidbClusterForPD()
+		tc.Spec.TiKV.StoreLabels = []string{"region", "zone", "rack"}
 		pmm, _, _, pdClient, podIndexer, nodeIndexer := newFakeTiKVMemberManager(tc)
 
 		if test.errWhenGetStores {

tests/actions.go

@@ -256,6 +256,7 @@ type TidbClusterConfig struct {
 	BlockWriteConfig blockwriter.Config
 	GrafanaClient    *metrics.Client
 	SubValues        string
+	TopologyKey      string
 }
 
 func (tc *TidbClusterConfig) String() string {
@@ -708,6 +709,13 @@ func (oa *operatorActions) CheckTidbClusterStatus(info *TidbClusterConfig) error
 			return false, nil
 		}
 
+		glog.V(4).Infof("check store labels")
+		if b, err := oa.storeLabelsIsSet(tc, info.TopologyKey); !b && err == nil {
+			return false, nil
+		} else if err != nil {
+			return false, err
+		}
+
 		glog.V(4).Infof("check tidb cluster begin passwordIsSet")
 		if b, err := oa.passwordIsSet(info); !b && err == nil {
 			return false, nil
@@ -1367,6 +1375,29 @@ func (oa *operatorActions) schedulerHAFn(tc *v1alpha1.TidbCluster) (bool, error)
 	return true, nil
 }
 
+func (oa *operatorActions) storeLabelsIsSet(tc *v1alpha1.TidbCluster, topologyKey string) (bool, error) {
+	pdCli := oa.pdControl.GetPDClient(tc)
+	for _, store := range tc.Status.TiKV.Stores {
+		storeID, err := strconv.ParseUint(store.ID, 10, 64)
+		if err != nil {
+			return false, err
+		}
+		storeInfo, err := pdCli.GetStore(storeID)
+		if err != nil {
+			return false, nil
+		}
+		if len(storeInfo.Store.Labels) == 0 {
+			return false, nil
+		}
+		for _, label := range storeInfo.Store.Labels {
+			if label.Key != topologyKey {
+				return false, nil
+			}
+		}
+	}
+	return true, nil
+}
+
 func (oa *operatorActions) passwordIsSet(clusterInfo *TidbClusterConfig) (bool, error) {
 	ns := clusterInfo.Namespace
 	tcName := clusterInfo.ClusterName

tests/cmd/e2e/main.go

@@ -73,6 +73,7 @@ func main() {
 	name1 := "e2e-cluster1"
 	name2 := "e2e-cluster2"
 	name3 := "e2e-pd-replicas-1"
+	topologyKey := "rack"
 	clusterInfos := []*tests.TidbClusterConfig{
 		{
 			Namespace:        name1,
@@ -110,7 +111,8 @@ func main() {
 				BatchSize:   1,
 				RawSize:     1,
 			},
-			SubValues:              tests.GetAffinityConfigOrDie(name1, name1),
+			TopologyKey:            topologyKey,
+			SubValues:              fmt.Sprintf("%s", tests.GetAffinityConfigOrDie(name1, name1, topologyKey, []string{topologyKey})),
 			EnableConfigMapRollout: true,
 			PDMaxReplicas:          3,
 			TiKVGrpcConcurrency:    4,
@@ -153,7 +155,8 @@ func main() {
 				BatchSize:   1,
 				RawSize:     1,
 			},
-			SubValues:              tests.GetAffinityConfigOrDie(name2, name2),
+			TopologyKey:            topologyKey,
+			SubValues:              fmt.Sprintf("%s", tests.GetAffinityConfigOrDie(name2, name2, topologyKey, []string{topologyKey})),
 			EnableConfigMapRollout: false,
 			PDMaxReplicas:          3,
 			TiKVGrpcConcurrency:    4,
@@ -176,7 +179,9 @@ func main() {
 				"pd.replicas":     "1",
 				"discovery.image": conf.OperatorImage,
 			},
-			SubValues: tests.GetAffinityConfigOrDie(name3, name2),
+
+			TopologyKey: topologyKey,
+			SubValues:   fmt.Sprintf("%s", tests.GetAffinityConfigOrDie(name3, name2, topologyKey, []string{topologyKey})),
 		},
 	}
 
@@ -335,7 +340,7 @@ func main() {
 	restoreClusterInfo.ClusterName = restoreClusterInfo.ClusterName + "-other"
 	restoreClusterInfo.InitSecretName = fmt.Sprintf("%s-set-secret", restoreClusterInfo.ClusterName)
 	restoreClusterInfo.BackupSecretName = fmt.Sprintf("%s-backup-secret", restoreClusterInfo.ClusterName)
-	restoreClusterInfo.SubValues = tests.GetAffinityConfigOrDie(restoreClusterInfo.ClusterName, restoreClusterInfo.Namespace)
+	restoreClusterInfo.SubValues = fmt.Sprintf("%s", tests.GetAffinityConfigOrDie(restoreClusterInfo.ClusterName, restoreClusterInfo.Namespace, topologyKey, []string{topologyKey}))
 
 	if err = oa.CleanTidbCluster(restoreClusterInfo); err != nil {
 		glog.Fatal(err)

tests/cmd/stability/stability.go

@@ -28,6 +28,7 @@ func newOperatorConfig() *tests.OperatorConfig {
 func newTidbClusterConfig(ns, clusterName string) *tests.TidbClusterConfig {
 	tidbVersion := cfg.GetTiDBVersionOrDie()
 
+	topologyKey := "rack"
 	return &tests.TidbClusterConfig{
 		Namespace:        ns,
 		ClusterName:      clusterName,
@@ -69,6 +70,6 @@ func newTidbClusterConfig(ns, clusterName string) *tests.TidbClusterConfig {
 		TiKVGrpcConcurrency: 4,
 		TiDBTokenLimit:      1000,
 		PDLogLevel:          "info",
-		SubValues:           tests.GetAffinityConfigOrDie(clusterName, ns),
+		SubValues:           tests.GetAffinityConfigOrDie(clusterName, ns, topologyKey, []string{topologyKey}),
 	}
 }

tests/util.go

@@ -89,6 +89,12 @@ func GetPodsByLabels(kubeCli kubernetes.Interface, node string, lables map[strin
 }
 
 var affinityTemp string = `{{.Kind}}:
+{{ $length := len .StoreLabels}} {{ if or (not .StoreLabels) (eq $length 0)}}
+{{else if eq .Kind "tikv"}}
+  storeLabels:
+{{range .StoreLabels}}  - {{.}}
+{{end}}
+{{end}}
   affinity:
     podAntiAffinity:
       preferredDuringSchedulingIgnoredDuringExecution:
@@ -98,7 +104,7 @@ var affinityTemp string = `{{.Kind}}:
             matchLabels:
               app.kubernetes.io/instance: {{.ClusterName}}
               app.kubernetes.io/component: {{.Kind}}
-          topologyKey: "rack"
+          topologyKey: {{.TopologyKey}}
           namespaces:
           - {{.Namespace}}
 `
@@ -108,26 +114,28 @@ type AffinityInfo struct {
 	Kind        string
 	Weight      int
 	Namespace   string
+	TopologyKey string
+	StoreLabels []string
 }
 
-func GetAffinityConfigOrDie(clusterName, namespace string) string {
+func GetAffinityConfigOrDie(clusterName, namespace, topologyKey string, storeLabels []string) string {
 	temp, err := template.New("dt-affinity").Parse(affinityTemp)
 	if err != nil {
 		slack.NotifyAndPanic(err)
 	}
 
 	pdbuff := new(bytes.Buffer)
-	err = temp.Execute(pdbuff, &AffinityInfo{ClusterName: clusterName, Kind: "pd", Weight: 50, Namespace: namespace})
+	err = temp.Execute(pdbuff, &AffinityInfo{ClusterName: clusterName, Kind: "pd", Weight: 50, Namespace: namespace, TopologyKey: topologyKey, StoreLabels: storeLabels})
 	if err != nil {
 		slack.NotifyAndPanic(err)
 	}
 	tikvbuff := new(bytes.Buffer)
-	err = temp.Execute(tikvbuff, &AffinityInfo{ClusterName: clusterName, Kind: "tikv", Weight: 50, Namespace: namespace})
+	err = temp.Execute(tikvbuff, &AffinityInfo{ClusterName: clusterName, Kind: "tikv", Weight: 50, Namespace: namespace, TopologyKey: topologyKey, StoreLabels: storeLabels})
 	if err != nil {
 		slack.NotifyAndPanic(err)
 	}
 	tidbbuff := new(bytes.Buffer)
-	err = temp.Execute(tidbbuff, &AffinityInfo{ClusterName: clusterName, Kind: "tidb", Weight: 50, Namespace: namespace})
+	err = temp.Execute(tidbbuff, &AffinityInfo{ClusterName: clusterName, Kind: "tidb", Weight: 50, Namespace: namespace, TopologyKey: topologyKey, StoreLabels: storeLabels})
 	if err != nil {
 		slack.NotifyAndPanic(err)
 	}