update store limit

Signed-off-by: okjiang <[email protected]>

Author: okJiang

server/cluster/cluster.go

@@ -2423,6 +2423,7 @@ func (c *RaftCluster) SetStoreLimit(storeID uint64, typ storelimit.Type, ratePer
 		log.Error("persist store limit meet error", errs.ZapError(err))
 		return err
 	}
+	c.refreshStoreRateLimit(storeID, typ)
 	log.Info("store limit changed", zap.Uint64("store-id", storeID), zap.String("type", typ.String()), zap.Float64("rate-per-min", ratePerMin))
 	return nil
 }
@@ -2441,10 +2442,35 @@ func (c *RaftCluster) SetAllStoresLimit(typ storelimit.Type, ratePerMin float64)
 		log.Error("persist store limit meet error", errs.ZapError(err))
 		return err
 	}
+	for storeID := range c.opt.GetAllStoresLimit() {
+		c.refreshStoreRateLimit(storeID, typ)
+	}
 	log.Info("all store limit changed", zap.String("type", typ.String()), zap.Float64("rate-per-min", ratePerMin))
 	return nil
 }
 
+// refreshStoreRateLimit applies the schedule config's store limit to the in-memory store limiter.
+func (c *RaftCluster) refreshStoreRateLimit(storeID uint64, limitType storelimit.Type) {
+	// Only v1 uses StoreRateLimit for AddPeer/RemovePeer.
+	if c.opt.GetStoreLimitVersion() != storelimit.VersionV1 {
+		return
+	}
+	store := c.GetStore(storeID)
+	if store == nil {
+		return
+	}
+	limit, ok := store.GetStoreLimit().(*storelimit.StoreRateLimit)
+	if !ok {
+		return
+	}
+	// Schedule config stores the unit in rate-per-minute, but limiter uses rate-per-second.
+	const storeBalanceBaseTime = float64(60)
+	ratePerSec := c.opt.GetStoreLimitByType(storeID, limitType) / storeBalanceBaseTime
+	if limit.Rate(limitType) != ratePerSec {
+		c.ResetStoreLimit(storeID, limitType, ratePerSec)
+	}
+}
+
 // SetAllStoresLimitTTL sets all store limit for a given type and rate with ttl.
 func (c *RaftCluster) SetAllStoresLimitTTL(typ storelimit.Type, ratePerMin float64, ttl time.Duration) error {
 	return c.opt.SetAllStoresLimitTTL(c.ctx, c.etcdClient, typ, ratePerMin, ttl)

server/cluster/cluster_test.go

@@ -2923,6 +2923,40 @@ func TestCheckCache(t *testing.T) {
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/schedule/checker/breakPatrol"))
 }
 
+func TestStoreLimitChangeRefreshLimiter(t *testing.T) {
+	re := require.New(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	_, opt, err := newTestScheduleConfig()
+	re.NoError(err)
+	// StoreRateLimit (v1) is used for AddPeer/RemovePeer and participates in scheduler filters.
+	opt.GetScheduleConfig().StoreLimitVersion = storelimit.VersionV1
+
+	rc := newTestRaftCluster(ctx, mockid.NewIDAllocator(), opt, storage.NewStorageWithMemoryBackend())
+
+	storeID := uint64(1)
+	store := core.NewStoreInfo(&metapb.Store{Id: storeID}, core.SetLastHeartbeatTS(time.Now()))
+	rc.PutStore(store)
+
+	// Simulate that the limiter has been set to an extremely low rate and already consumed,
+	// which would make scheduler filters throttle this store.
+	lowRatePerSec := float64(0.0001) / 60.0
+	rc.ResetStoreLimit(storeID, storelimit.AddPeer, lowRatePerSec)
+	store = rc.GetStore(storeID)
+	re.NotNil(store)
+	re.True(store.GetStoreLimit().Take(storelimit.RegionInfluence[storelimit.AddPeer], storelimit.AddPeer, constant.Low))
+	re.False(store.IsAvailable(storelimit.AddPeer, constant.Low))
+
+	// Increase store limit via config API. Without refreshing the in-memory limiter,
+	// the store would remain throttled and be filtered out.
+	re.NoError(rc.SetStoreLimit(storeID, storelimit.AddPeer, 30))
+	store = rc.GetStore(storeID)
+	re.NotNil(store)
+	re.True(store.IsAvailable(storelimit.AddPeer, constant.Low))
+}
+
 func TestPatrolRegionConcurrency(t *testing.T) {
 	re := require.New(t)