diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index bd09813f01ea..05c7008b7e85 100644 --- a/server/cluster/cluster.go +++ b/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) } diff --git a/server/config/persist_options.go b/server/config/persist_options.go index 2aa793d2f23f..f13af0a4c321 100644 --- a/server/config/persist_options.go +++ b/server/config/persist_options.go @@ -207,14 +207,12 @@ func (o *PersistOptions) SetAllStoresLimit(typ storelimit.Type, ratePerMin float switch typ { case storelimit.AddPeer: DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, ratePerMin) - DefaultTiFlashStoreLimit.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) - DefaultTiFlashStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, ratePerMin) for storeID := range v.StoreLimit { sc := StoreLimitConfig{AddPeer: v.StoreLimit[storeID].AddPeer, RemovePeer: ratePerMin} v.StoreLimit[storeID] = sc diff --git a/server/core/store.go b/server/core/store.go index 48b986ef9f8e..9642650ab79f 100644 --- a/server/core/store.go +++ b/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 +} diff --git a/server/grpc_service.go b/server/grpc_service.go index 1914c25de60e..0d0504db66da 100644 --- a/server/grpc_service.go +++ b/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(), diff --git a/server/server.go b/server/server.go index 488253f7a339..543a60f49cba 100644 --- a/server/server.go +++ b/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") } } diff --git a/server/util.go b/server/util.go index 45e11a706fff..6da2aad3cbe9 100644 --- a/server/util.go +++ b/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 -}