schedule: fix the zero store limit (#2564)

Signed-off-by: Ryan Leung <rleungx@gmail.com>

server/cluster/cluster.go

@@ -1684,18 +1684,19 @@ func (c *RaftCluster) GetAllStoresLimit() map[uint64]config.StoreLimitConfig {
 }
 
 // AddStoreLimit add a store limit for a given store ID.
-func (c *RaftCluster) AddStoreLimit(storeID uint64, isTiFlashStore bool) {
+func (c *RaftCluster) AddStoreLimit(store *metapb.Store) {
 	cfg := c.opt.GetScheduleConfig().Clone()
 	sc := config.StoreLimitConfig{
 		AddPeer:    config.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer),
 		RemovePeer: config.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer),
 	}
-	if isTiFlashStore {
+	if core.IsTiFlashStore(store) {
 		sc = config.StoreLimitConfig{
 			AddPeer:    config.DefaultTiFlashStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer),
 			RemovePeer: config.DefaultTiFlashStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer),
 		}
 	}
+	storeID := store.GetId()
 	cfg.StoreLimit[storeID] = sc
 	c.opt.SetScheduleConfig(cfg)
 }

server/config/config.go

@@ -226,7 +226,7 @@ var (
 
 // StoreLimit is the default limit of adding peer and removing peer when putting stores.
 type StoreLimit struct {
-	mu sync.Mutex
+	mu sync.RWMutex
 	// AddPeer is the default rate of adding peers for store limit (per minute).
 	AddPeer float64
 	// RemovePeer is the default rate of removing peers for store limit (per minute).
@@ -247,8 +247,8 @@ func (sl *StoreLimit) SetDefaultStoreLimit(typ storelimit.Type, ratePerMin float
 
 // GetDefaultStoreLimit gets the default store limit for a given type.
 func (sl *StoreLimit) GetDefaultStoreLimit(typ storelimit.Type) float64 {
-	sl.mu.Lock()
-	defer sl.mu.Unlock()
+	sl.mu.RLock()
+	defer sl.mu.RUnlock()
 	switch typ {
 	case storelimit.AddPeer:
 		return sl.AddPeer

server/config/persist_options.go

@@ -206,11 +206,13 @@ func (o *PersistOptions) SetAllStoresLimit(typ storelimit.Type, ratePerMin float
 	v := o.GetScheduleConfig().Clone()
 	switch typ {
 	case storelimit.AddPeer:
+		DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, ratePerMin)
 		for storeID := range v.StoreLimit {
 			sc := StoreLimitConfig{AddPeer: ratePerMin, RemovePeer: v.StoreLimit[storeID].RemovePeer}
 			v.StoreLimit[storeID] = sc
 		}
 	case storelimit.RemovePeer:
+		DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, ratePerMin)
 		for storeID := range v.StoreLimit {
 			sc := StoreLimitConfig{AddPeer: v.StoreLimit[storeID].AddPeer, RemovePeer: ratePerMin}
 			v.StoreLimit[storeID] = sc
@@ -270,7 +272,14 @@ func (o *PersistOptions) GetStoreLimit(storeID uint64) StoreLimitConfig {
 	if limit, ok := o.GetScheduleConfig().StoreLimit[storeID]; ok {
 		return limit
 	}
-	return StoreLimitConfig{0, 0}
+	cfg := o.GetScheduleConfig().Clone()
+	sc := StoreLimitConfig{
+		AddPeer:    DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer),
+		RemovePeer: DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer),
+	}
+	cfg.StoreLimit[storeID] = sc
+	o.SetScheduleConfig(cfg)
+	return o.GetScheduleConfig().StoreLimit[storeID]
 }
 
 // GetStoreLimitByType returns the limit of a store with a given type.

server/core/store.go

@@ -671,3 +671,14 @@ func (s *StoresInfo) UpdateStoreStatus(storeID uint64, leaderCount int, regionCo
 		s.SetStore(newStore)
 	}
 }
+
+// IsTiFlashStore used to judge flash store.
+// FIXME: remove the hack way
+func IsTiFlashStore(store *metapb.Store) bool {
+	for _, l := range store.GetLabels() {
+		if l.GetKey() == "engine" && l.GetValue() == "tiflash" {
+			return true
+		}
+	}
+	return false
+}

server/grpc_service.go

@@ -224,7 +224,7 @@ func (s *Server) PutStore(ctx context.Context, request *pdpb.PutStoreRequest) (*
 	}
 
 	// NOTE: can be removed when placement rules feature is enabled by default.
-	if !s.GetConfig().Replication.EnablePlacementRules && isTiFlashStore(store) {
+	if !s.GetConfig().Replication.EnablePlacementRules && core.IsTiFlashStore(store) {
 		return nil, status.Errorf(codes.FailedPrecondition, "placement rules is disabled")
 	}
 
@@ -235,11 +235,7 @@ func (s *Server) PutStore(ctx context.Context, request *pdpb.PutStoreRequest) (*
 	log.Info("put store ok", zap.Stringer("store", store))
 	rc.OnStoreVersionChange()
 	CheckPDVersion(s.persistOptions)
-	if isTiFlashStore(store) {
-		rc.AddStoreLimit(store.GetId(), true /* isTiFlashStore*/)
-	} else {
-		rc.AddStoreLimit(store.GetId(), false /* isTiFlashStore*/)
-	}
+	rc.AddStoreLimit(store)
 
 	return &pdpb.PutStoreResponse{
 		Header:            s.header(),

server/server.go

@@ -765,7 +765,7 @@ func (s *Server) SetReplicationConfig(cfg config.ReplicationConfig) error {
 		} else {
 			// NOTE: can be removed after placement rules feature is enabled by default.
 			for _, s := range raftCluster.GetStores() {
-				if !s.IsTombstone() && isTiFlashStore(s.GetMeta()) {
+				if !s.IsTombstone() && core.IsTiFlashStore(s.GetMeta()) {
 					return errors.New("cannot disable placement rules with TiFlash nodes")
 				}
 			}

server/util.go

@@ -20,7 +20,6 @@ import (
 	"path"
 	"time"
 
-	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/log"
 	"github.com/pingcap/pd/v4/pkg/etcdutil"
@@ -182,14 +181,3 @@ func checkBootstrapRequest(clusterID uint64, req *pdpb.BootstrapRequest) error {
 
 	return nil
 }
-
-// isTiFlashStore used to judge flash store.
-// FIXME: remove the hack way
-func isTiFlashStore(store *metapb.Store) bool {
-	for _, l := range store.GetLabels() {
-		if l.GetKey() == "engine" && l.GetValue() == "tiflash" {
-			return true
-		}
-	}
-	return false
-}

tests/server/cluster/cluster_test.go

@@ -960,3 +960,58 @@ func (s *clusterTestSuite) TestOfflineStoreLimit(c *C) {
 		c.Assert(oc.RemoveOperator(op), IsTrue)
 	}
 }
+
+func (s *clusterTestSuite) TestUpgradeStoreLimit(c *C) {
+	tc, err := tests.NewTestCluster(s.ctx, 1)
+	defer tc.Destroy()
+	c.Assert(err, IsNil)
+	err = tc.RunInitialServers()
+	c.Assert(err, IsNil)
+	tc.WaitLeader()
+	leaderServer := tc.GetServer(tc.GetLeader())
+	grpcPDClient := testutil.MustNewGrpcClient(c, leaderServer.GetAddr())
+	clusterID := leaderServer.GetClusterID()
+	bootstrapCluster(c, clusterID, grpcPDClient, "127.0.0.1:0")
+	rc := leaderServer.GetRaftCluster()
+	c.Assert(rc, NotNil)
+	rc.SetStorage(core.NewStorage(kv.NewMemoryKV()))
+	store := newMetaStore(1, "127.0.1.1:0", "4.0.0", metapb.StoreState_Up)
+	_, err = putStore(c, grpcPDClient, clusterID, store)
+	c.Assert(err, IsNil)
+	r := &metapb.Region{
+		Id: 1,
+		RegionEpoch: &metapb.RegionEpoch{
+			ConfVer: 1,
+			Version: 1,
+		},
+		StartKey: []byte{byte(2)},
+		EndKey:   []byte{byte(3)},
+		Peers:    []*metapb.Peer{{Id: 11, StoreId: uint64(1)}},
+	}
+	region := core.NewRegionInfo(r, r.Peers[0], core.SetApproximateSize(10))
+
+	err = rc.HandleRegionHeartbeat(region)
+	c.Assert(err, IsNil)
+
+	// restart PD
+	// Here we use an empty storelimit to simulate the upgrade progress.
+	opt := rc.GetOpt()
+	scheduleCfg := opt.GetScheduleConfig()
+	scheduleCfg.StoreLimit = map[uint64]config.StoreLimitConfig{}
+	c.Assert(leaderServer.GetServer().SetScheduleConfig(*scheduleCfg), IsNil)
+	err = leaderServer.Stop()
+	c.Assert(err, IsNil)
+	err = leaderServer.Run()
+	c.Assert(err, IsNil)
+
+	oc := rc.GetOperatorController()
+	// only can add 5 remove peer operators on store 1
+	for i := uint64(1); i <= 5; i++ {
+		op := operator.NewOperator("test", "test", 1, &metapb.RegionEpoch{ConfVer: 1, Version: 1}, operator.OpRegion, operator.RemovePeer{FromStore: 1})
+		c.Assert(oc.AddOperator(op), IsTrue)
+		c.Assert(oc.RemoveOperator(op), IsTrue)
+	}
+	op := operator.NewOperator("test", "test", 1, &metapb.RegionEpoch{ConfVer: 1, Version: 1}, operator.OpRegion, operator.RemovePeer{FromStore: 1})
+	c.Assert(oc.AddOperator(op), IsFalse)
+	c.Assert(oc.RemoveOperator(op), IsFalse)
+}