From a3b0e6c14c7a762a800c8348de251fd1d4ededaf Mon Sep 17 00:00:00 2001
From: zyy <zhangyunyu@zhihu.com>
Date: Thu, 25 Feb 2021 16:16:03 +0800
Subject: [PATCH] add store labels (#3784)

* add store labels

* Update pkg/apis/pingcap/v1alpha1/types.go

Co-authored-by: Lonng <heng@lonng.org>

* Update pkg/manager/member/tikv_member_manager.go

Co-authored-by: Lonng <heng@lonng.org>

* Update pkg/apis/pingcap/v1alpha1/types.go

Co-authored-by: Lonng <heng@lonng.org>

* fix var

* fix

* generate code

Co-authored-by: Lonng <heng@lonng.org>
Co-authored-by: DanielZhangQD <36026334+DanielZhangQD@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 c40614dd118..c27d0dee2f8 100644
--- a/docs/api-references/docs.md
+++ b/docs/api-references/docs.md
@@ -18057,6 +18057,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 b7623926ee4..fc7e7442255 100644
--- a/manifests/crd.yaml
+++ b/manifests/crd.yaml
@@ -12788,6 +12788,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 7a8aef4ded0..aa7036cdb6b 100644
--- a/pkg/apis/pingcap/v1alpha1/openapi_generated.go
+++ b/pkg/apis/pingcap/v1alpha1/openapi_generated.go
@@ -9803,6 +9803,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 bdb9964ac1e..d6f7d5666f5 100644
--- a/pkg/apis/pingcap/v1alpha1/types.go
+++ b/pkg/apis/pingcap/v1alpha1/types.go
@@ -468,6 +468,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 b5f784687c2..130bee5f3d9 100644
--- a/pkg/apis/pingcap/v1alpha1/zz_generated.deepcopy.go
+++ b/pkg/apis/pingcap/v1alpha1/zz_generated.deepcopy.go
@@ -7204,6 +7204,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 b831ee4062d..6102ce6550f 100644
--- a/pkg/manager/member/tikv_member_manager.go
+++ b/pkg/manager/member/tikv_member_manager.go
@@ -816,8 +816,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
 	}
 
@@ -843,7 +843,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