*: add option to check label strictly

server/api/label_test.go

@@ -14,14 +14,18 @@
 package api
 
 import (
+	"context"
 	"fmt"
+	"strings"
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/pd/server"
 )
 
 var _ = Suite(&testLabelsStoreSuite{})
+var _ = Suite(&testStrictlyLabelsStoreSuite{})
 
 type testLabelsStoreSuite struct {
 	svr       *server.Server
@@ -102,7 +106,7 @@ func (s *testLabelsStoreSuite) SetUpSuite(c *C) {
 		},
 	}
 
-	s.svr, s.cleanup = mustNewServer(c)
+	s.svr, s.cleanup = mustNewServer(c, func(cfg *server.Config) { cfg.Replication.StrictlyMatchLabel = false })
 	mustWaitLeader(c, []*server.Server{s.svr})
 
 	addr := s.svr.GetAddr()
@@ -170,3 +174,107 @@ func (s *testLabelsStoreSuite) TestStoresLabelFilter(c *C) {
 	_, err := newStoresLabelFilter("test", ".[test")
 	c.Assert(err, NotNil)
 }
+
+type testStrictlyLabelsStoreSuite struct {
+	svr       *server.Server
+	cleanup   cleanUpFunc
+	urlPrefix string
+}
+
+func (s *testStrictlyLabelsStoreSuite) SetUpSuite(c *C) {
+	s.svr, s.cleanup = mustNewServer(c, func(cfg *server.Config) {
+		cfg.Replication.LocationLabels = []string{"zone", "disk"}
+		cfg.Replication.StrictlyMatchLabel = true
+	})
+	mustWaitLeader(c, []*server.Server{s.svr})
+
+	addr := s.svr.GetAddr()
+	s.urlPrefix = fmt.Sprintf("%s%s/api/v1", addr, apiPrefix)
+
+	mustBootstrapCluster(c, s.svr)
+}
+
+func (s *testStrictlyLabelsStoreSuite) TestStoreMatch(c *C) {
+	cases := []struct {
+		store       *metapb.Store
+		valid       bool
+		expectError string
+	}{
+		{
+			store: &metapb.Store{
+				Id:      1,
+				Address: "tikv1",
+				State:   metapb.StoreState_Up,
+				Labels: []*metapb.StoreLabel{
+					{
+						Key:   "zone",
+						Value: "us-west-1",
+					},
+					{
+						Key:   "disk",
+						Value: "ssd",
+					},
+				},
+				Version: "3.0.0",
+			},
+			valid: true,
+		},
+		{
+			store: &metapb.Store{
+				Id:      2,
+				Address: "tikv2",
+				State:   metapb.StoreState_Up,
+				Labels:  []*metapb.StoreLabel{},
+				Version: "3.0.0",
+			},
+			valid:       false,
+			expectError: "label configuration is incorrect",
+		},
+		{
+			store: &metapb.Store{
+				Id:      2,
+				Address: "tikv2",
+				State:   metapb.StoreState_Up,
+				Labels: []*metapb.StoreLabel{
+					{
+						Key:   "zone",
+						Value: "cn-beijing-1",
+					},
+					{
+						Key:   "disk",
+						Value: "ssd",
+					},
+					{
+						Key:   "other",
+						Value: "unknown",
+					},
+				},
+				Version: "3.0.0",
+			},
+			valid:       false,
+			expectError: "key matching the label was not found",
+		},
+	}
+
+	for _, t := range cases {
+		_, err := s.svr.PutStore(context.Background(), &pdpb.PutStoreRequest{
+			Header: &pdpb.RequestHeader{ClusterId: s.svr.ClusterID()},
+			Store: &metapb.Store{
+				Id:      t.store.Id,
+				Address: fmt.Sprintf("tikv%d", t.store.Id),
+				State:   t.store.State,
+				Labels:  t.store.Labels,
+				Version: t.store.Version,
+			},
+		})
+		if t.valid {
+			c.Assert(err, IsNil)
+		} else {
+			c.Assert(strings.Contains(err.Error(), t.expectError), IsTrue)
+		}
+	}
+}
+
+func (s *testStrictlyLabelsStoreSuite) TearDownSuite(c *C) {
+	s.cleanup()
+}

server/api/server_test.go

@@ -70,14 +70,14 @@ func cleanServer(cfg *server.Config) {
 
 type cleanUpFunc func()
 
-func mustNewServer(c *C) (*server.Server, cleanUpFunc) {
-	_, svrs, cleanup := mustNewCluster(c, 1)
+func mustNewServer(c *C, opts ...func(cfg *server.Config)) (*server.Server, cleanUpFunc) {
+	_, svrs, cleanup := mustNewCluster(c, 1, opts...)
 	return svrs[0], cleanup
 }
 
 var zapLogOnce sync.Once
 
-func mustNewCluster(c *C, num int) ([]*server.Config, []*server.Server, cleanUpFunc) {
+func mustNewCluster(c *C, num int, opts ...func(cfg *server.Config)) ([]*server.Config, []*server.Server, cleanUpFunc) {
 	svrs := make([]*server.Server, 0, num)
 	cfgs := server.NewTestMultiConfig(c, num)
 
@@ -89,6 +89,9 @@ func mustNewCluster(c *C, num int) ([]*server.Config, []*server.Server, cleanUpF
 			zapLogOnce.Do(func() {
 				log.ReplaceGlobals(cfg.GetZapLogger(), cfg.GetZapLogProperties())
 			})
+			for _, opt := range opts {
+				opt(cfg)
+			}
 			s, err := server.CreateServer(cfg, NewHandler)
 			c.Assert(err, IsNil)
 			err = s.Run(context.TODO())

server/api/store_test.go

@@ -19,6 +19,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"
 
 	. "github.com/pingcap/check"
@@ -151,6 +152,12 @@ func (s *testStoreSuite) TestStoreLabel(c *C) {
 	b, err := json.Marshal(labels)
 	c.Assert(err, IsNil)
 	err = postJSON(url+"/label", b)
+	c.Assert(strings.Contains(err.Error(), "key matching the label was not found"), IsTrue)
+	locationLabels := map[string]string{"location-labels": "zone,host"}
+	ll, _ := json.Marshal(locationLabels)
+	err = postJSON(s.urlPrefix+"/config", ll)
+	c.Assert(err, IsNil)
+	err = postJSON(url+"/label", b)
 	c.Assert(err, IsNil)
 
 	err = readJSONWithURL(url, &info)
@@ -161,6 +168,12 @@ func (s *testStoreSuite) TestStoreLabel(c *C) {
 	}
 
 	// Test merge.
+	// disable label match check.
+	labelCheck := map[string]string{"strictly-match-label": "false"}
+	lc, _ := json.Marshal(labelCheck)
+	err = postJSON(s.urlPrefix+"/config", lc)
+	c.Assert(err, IsNil)
+
 	labels = map[string]string{"zack": "zack1", "Host": "host1"}
 	b, err = json.Marshal(labels)
 	c.Assert(err, IsNil)

server/cluster.go

@@ -402,11 +402,27 @@ func (c *RaftCluster) putStore(store *metapb.Store) error {
 		)
 	}
 	// Check location labels.
-	for _, k := range c.cachedCluster.GetLocationLabels() {
+	keysSet := make(map[string]struct{})
+	for _, k := range cluster.GetLocationLabels() {
+		keysSet[k] = struct{}{}
 		if v := s.GetLabelValue(k); len(v) == 0 {
-			log.Warn("missing location label",
+			log.Warn("label configuration is incorrect",
 				zap.Stringer("store", s.GetMeta()),
 				zap.String("label-key", k))
+			if cluster.GetStrictlyMatchLabel() {
+				return errors.Errorf("label configuration is incorrect, need to specify the key: %s ", k)
+			}
+		}
+	}
+	for _, label := range s.GetLabels() {
+		key := label.GetKey()
+		if _, ok := keysSet[key]; !ok {
+			log.Warn("not found the key match with the store label",
+				zap.Stringer("store", s.GetMeta()),
+				zap.String("label-key", key))
+			if cluster.GetStrictlyMatchLabel() {
+				return errors.Errorf("key matching the label was not found in the PD, store label key: %s ", key)
+			}
 		}
 	}
 	return cluster.putStore(s)

server/cluster_info.go

@@ -753,6 +753,10 @@ func (c *clusterInfo) GetLocationLabels() []string {
 	return c.opt.GetLocationLabels()
 }
 
+func (c *clusterInfo) GetStrictlyMatchLabel() bool {
+	return c.opt.rep.GetStrictlyMatchLabel()
+}
+
 func (c *clusterInfo) GetHotRegionCacheHitsThreshold() int {
 	return c.opt.GetHotRegionCacheHitsThreshold()
 }

server/config.go

@@ -197,7 +197,8 @@ const (
 
 	defaultLeaderPriorityCheckInterval = time.Minute
 
-	defaultUseRegionStorage = true
+	defaultUseRegionStorage   = true
+	defaultStrictlyMatchLabel = true
 )
 
 func adjustString(v *string, defValue string) {
@@ -681,14 +682,17 @@ type ReplicationConfig struct {
 	// 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.
 	LocationLabels typeutil.StringSlice `toml:"location-labels,omitempty" json:"location-labels"`
+	// StrictlyMatchLabel strictly checks if the label of TiKV is matched with LocaltionLabels.
+	StrictlyMatchLabel bool `toml:"strictly-match-label,omitempty" json:"strictly-match-label,string"`
 }
 
 func (c *ReplicationConfig) clone() *ReplicationConfig {
 	locationLabels := make(typeutil.StringSlice, len(c.LocationLabels))
 	copy(locationLabels, c.LocationLabels)
 	return &ReplicationConfig{
-		MaxReplicas:    c.MaxReplicas,
-		LocationLabels: locationLabels,
+		MaxReplicas:        c.MaxReplicas,
+		LocationLabels:     locationLabels,
+		StrictlyMatchLabel: c.StrictlyMatchLabel,
 	}
 }
 
@@ -704,6 +708,9 @@ func (c *ReplicationConfig) validate() error {
 
 func (c *ReplicationConfig) adjust(meta *configMetaData) error {
 	adjustUint64(&c.MaxReplicas, defaultMaxReplicas)
+	if !meta.IsDefined("strictly-match-label") {
+		c.StrictlyMatchLabel = defaultStrictlyMatchLabel
+	}
 	return c.validate()
 }
 

server/option.go

@@ -430,6 +430,11 @@ func (r *Replication) GetLocationLabels() []string {
 	return r.load().LocationLabels
 }
 
+// GetStrictlyMatchLabel returns whether check label strict.
+func (r *Replication) GetStrictlyMatchLabel() bool {
+	return r.load().StrictlyMatchLabel
+}
+
 // namespaceOption is a wrapper to access the configuration safely.
 type namespaceOption struct {
 	namespaceCfg atomic.Value

server/schedule/mockcluster.go

@@ -511,6 +511,7 @@ const (
 	defaultLowSpaceRatio               = 0.8
 	defaultHighSpaceRatio              = 0.6
 	defaultHotRegionCacheHitsThreshold = 3
+	defaultStrictlyMatchLabel          = true
 )
 
 // MockSchedulerOptions is a mock of SchedulerOptions
@@ -530,6 +531,7 @@ type MockSchedulerOptions struct {
 	MaxStoreDownTime             time.Duration
 	MaxReplicas                  int
 	LocationLabels               []string
+	StrictlyMatchLabel           bool
 	HotRegionCacheHitsThreshold  int
 	TolerantSizeRatio            float64
 	LowSpaceRatio                float64
@@ -559,6 +561,7 @@ func NewMockSchedulerOptions() *MockSchedulerOptions {
 	mso.SplitMergeInterval = defaultSplitMergeInterval
 	mso.MaxStoreDownTime = defaultMaxStoreDownTime
 	mso.MaxReplicas = defaultMaxReplicas
+	mso.StrictlyMatchLabel = defaultStrictlyMatchLabel
 	mso.HotRegionCacheHitsThreshold = defaultHotRegionCacheHitsThreshold
 	mso.MaxPendingPeerCount = defaultMaxPendingPeerCount
 	mso.TolerantSizeRatio = defaultTolerantSizeRatio
@@ -637,6 +640,11 @@ func (mso *MockSchedulerOptions) GetLocationLabels() []string {
 	return mso.LocationLabels
 }
 
+// GetStrictlyMatchLabel mock method
+func (mso *MockSchedulerOptions) GetStrictlyMatchLabel() bool {
+	return mso.StrictlyMatchLabel
+}
+
 // GetHotRegionCacheHitsThreshold mock method
 func (mso *MockSchedulerOptions) GetHotRegionCacheHitsThreshold() int {
 	return mso.HotRegionCacheHitsThreshold

server/schedule/opts.go

@@ -43,6 +43,7 @@ type Options interface {
 
 	GetMaxReplicas() int
 	GetLocationLabels() []string
+	GetStrictlyMatchLabel() bool
 
 	GetHotRegionCacheHitsThreshold() int
 	GetTolerantSizeRatio() float64

tests/pdctl/label/label_test.go

@@ -41,7 +41,7 @@ func (s *labelTestSuite) SetUpSuite(c *C) {
 func (s *labelTestSuite) TestLabel(c *C) {
 	c.Parallel()
 
-	cluster, err := tests.NewTestCluster(1)
+	cluster, err := tests.NewTestCluster(1, func(cfg *server.Config) { cfg.Replication.StrictlyMatchLabel = false })
 	c.Assert(err, IsNil)
 	err = cluster.RunInitialServers()
 	c.Assert(err, IsNil)