From 6ebf6def428315dd06ea8fc899371e0e0ba6b2b3 Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Thu, 25 Feb 2021 17:26:03 +0800
Subject: [PATCH] add store labels (#3784) (#3813)

* cherry pick #3784 to release-1.1

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>

* resolve conflict

Co-authored-by: zyy <zhangyunyu@zhihu.com>
Co-authored-by: DanielZhangQD <zhanghailong810@aliyun.com>
Co-authored-by: Ti Chi Robot <71242396+ti-chi-bot@users.noreply.github.com>
---
 docs/api-references/docs.md                        | 12 ++++++++++++
 manifests/crd.yaml                                 |  4 ++++
 pkg/apis/pingcap/v1alpha1/openapi_generated.go     | 14 ++++++++++++++
 pkg/apis/pingcap/v1alpha1/types.go                 |  4 ++++
 pkg/apis/pingcap/v1alpha1/zz_generated.deepcopy.go |  5 +++++
 pkg/manager/member/tikv_member_manager.go          |  6 +++---
 6 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/docs/api-references/docs.md b/docs/api-references/docs.md
index a42cce62f6e..6eb86928148 100644
--- a/docs/api-references/docs.md
+++ b/docs/api-references/docs.md
@@ -16395,6 +16395,18 @@ Defaults to 10m</p>
 <p>StorageVolumes configure additional storage for TiKV pods.</p>
 </td>
 </tr>
+<tr>
+<td>
+<code>storeLabels</code></br>
+<em>
+[]string
+</em>
+</td>
+<td>
+<em>(Optional)</em>
+<p>StoreLabels configures additional labels for TiKV stores.</p>
+</td>
+</tr>
 </tbody>
 </table>
 <h3 id="tikvstatus">TiKVStatus</h3>
diff --git a/manifests/crd.yaml b/manifests/crd.yaml
index 7b883ebd7fe..9f4d9ef7dd0 100644
--- a/manifests/crd.yaml
+++ b/manifests/crd.yaml
@@ -12784,6 +12784,10 @@ spec:
                 storageVolumes:
                   items: {}
                   type: array
+                storeLabels:
+                  items:
+                    type: string
+                  type: array
                 terminationGracePeriodSeconds:
                   format: int64
                   type: integer
diff --git a/pkg/apis/pingcap/v1alpha1/openapi_generated.go b/pkg/apis/pingcap/v1alpha1/openapi_generated.go
index 36c6bd7d08b..5aa8736d9d7 100644
--- a/pkg/apis/pingcap/v1alpha1/openapi_generated.go
+++ b/pkg/apis/pingcap/v1alpha1/openapi_generated.go
@@ -9176,6 +9176,20 @@ func schema_pkg_apis_pingcap_v1alpha1_TiKVSpec(ref common.ReferenceCallback) com
 							},
 						},
 					},
+					"storeLabels": {
+						SchemaProps: spec.SchemaProps{
+							Description: "StoreLabels configures additional labels for TiKV stores.",
+							Type:        []string{"array"},
+							Items: &spec.SchemaOrArray{
+								Schema: &spec.Schema{
+									SchemaProps: spec.SchemaProps{
+										Type:   []string{"string"},
+										Format: "",
+									},
+								},
+							},
+						},
+					},
 				},
 				Required: []string{"replicas"},
 			},
diff --git a/pkg/apis/pingcap/v1alpha1/types.go b/pkg/apis/pingcap/v1alpha1/types.go
index c87ff59add9..8a08ebb22ee 100644
--- a/pkg/apis/pingcap/v1alpha1/types.go
+++ b/pkg/apis/pingcap/v1alpha1/types.go
@@ -450,6 +450,10 @@ type TiKVSpec struct {
 	// StorageVolumes configure additional storage for TiKV pods.
 	// +optional
 	StorageVolumes []StorageVolume `json:"storageVolumes,omitempty"`
+
+	// StoreLabels configures additional labels for TiKV stores.
+	// +optional
+	StoreLabels []string `json:"storeLabels,omitempty"`
 }
 
 // TiFlashSpec contains details of TiFlash members
diff --git a/pkg/apis/pingcap/v1alpha1/zz_generated.deepcopy.go b/pkg/apis/pingcap/v1alpha1/zz_generated.deepcopy.go
index 731fdf04f12..d548329f034 100644
--- a/pkg/apis/pingcap/v1alpha1/zz_generated.deepcopy.go
+++ b/pkg/apis/pingcap/v1alpha1/zz_generated.deepcopy.go
@@ -6701,6 +6701,11 @@ func (in *TiKVSpec) DeepCopyInto(out *TiKVSpec) {
 			(*in)[i].DeepCopyInto(&(*out)[i])
 		}
 	}
+	if in.StoreLabels != nil {
+		in, out := &in.StoreLabels, &out.StoreLabels
+		*out = make([]string, len(*in))
+		copy(*out, *in)
+	}
 	return
 }
 
diff --git a/pkg/manager/member/tikv_member_manager.go b/pkg/manager/member/tikv_member_manager.go
index c9df08ade6a..f4130e6966c 100644
--- a/pkg/manager/member/tikv_member_manager.go
+++ b/pkg/manager/member/tikv_member_manager.go
@@ -807,8 +807,8 @@ func (m *tikvMemberManager) setStoreLabelsForTiKV(tc *v1alpha1.TidbCluster) (int
 		return setCount, err
 	}
 
-	locationLabels := []string(config.Replication.LocationLabels)
-	if locationLabels == nil {
+	storeLabels := append(config.Replication.LocationLabels, tc.Spec.TiKV.StoreLabels...)
+	if storeLabels == nil {
 		return setCount, nil
 	}
 
@@ -834,7 +834,7 @@ func (m *tikvMemberManager) setStoreLabelsForTiKV(tc *v1alpha1.TidbCluster) (int
 		}
 
 		nodeName := pod.Spec.NodeName
-		ls, err := m.getNodeLabels(nodeName, locationLabels)
+		ls, err := m.getNodeLabels(nodeName, storeLabels)
 		if err != nil || len(ls) == 0 {
 			klog.Warningf("node: [%s] has no node labels, skipping set store labels for Pod: [%s/%s]", nodeName, ns, podName)
 			continue