don't limit the scatter region

Signed-off-by: Ryan Leung <rleungx@gmail.com>
tikv · Jun 13, 2023 · fd8820a · fd8820a
1 parent 170d287
commit fd8820a
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 1 deletion.
diff --git a/pkg/schedule/operator/builder.go b/pkg/schedule/operator/builder.go
@@ -792,7 +792,11 @@ func (b *Builder) execRemovePeer(peer *metapb.Peer) {
 	if store != nil {
 		isDownStore = store.DownTime() > b.GetOpts().GetMaxStoreDownTime()
 	}
-	b.steps = append(b.steps, RemovePeer{FromStore: removeStoreID, PeerID: peer.GetId(), IsDownStore: isDownStore})
+	if b.lightWeight {
+		b.steps = append(b.steps, RemovePeer{FromStore: removeStoreID, PeerID: peer.GetId(), IsDownStore: isDownStore, IsLightWeight: b.lightWeight})
+	} else {
+		b.steps = append(b.steps, RemovePeer{FromStore: removeStoreID, PeerID: peer.GetId(), IsDownStore: isDownStore})
+	}
 	delete(b.currentPeers, removeStoreID)
 	delete(b.toRemove, removeStoreID)
 }

diff --git a/pkg/schedule/operator/step.go b/pkg/schedule/operator/step.go
@@ -588,6 +588,7 @@ func (pl PromoteLearner) GetCmd(_ *core.RegionInfo, useConfChangeV2 bool) *pdpb.
 // RemovePeer is an OpStep that removes a region peer.
 type RemovePeer struct {
 	FromStore, PeerID uint64
+	IsLightWeight     bool
 	IsDownStore       bool
 }
 
@@ -635,6 +636,10 @@ func (rp RemovePeer) Influence(opInfluence OpInfluence, region *core.RegionInfo)
 		return
 	}
 
+	if rp.IsLightWeight {
+		return
+	}
+
 	if rp.IsDownStore && regionSize > storelimit.SmallRegionThreshold {
 		regionSize = storelimit.SmallRegionThreshold
 	}

diff --git a/pkg/schedule/scatter/region_scatterer_test.go b/pkg/schedule/scatter/region_scatterer_test.go
@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/pd/pkg/core"
+	"github.com/tikv/pd/pkg/core/storelimit"
 	"github.com/tikv/pd/pkg/mock/mockcluster"
 	"github.com/tikv/pd/pkg/mock/mockconfig"
 	"github.com/tikv/pd/pkg/schedule/hbstream"
@@ -781,3 +782,37 @@ func isPeerCountChanged(op *operator.Operator) bool {
 	}
 	return add != remove
 }
+
+func TestRemoveStoreLimit(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	opt := mockconfig.NewTestOptions()
+	tc := mockcluster.NewCluster(ctx, opt)
+	stream := hbstream.NewTestHeartbeatStreams(ctx, tc.ID, tc, false)
+	oc := operator.NewController(ctx, tc.GetBasicCluster(), tc.GetOpts(), stream)
+
+	// Add stores 1~6.
+	for i := uint64(1); i <= 5; i++ {
+		tc.AddRegionStore(i, 0)
+		tc.SetStoreLimit(i, storelimit.AddPeer, 1)
+		tc.SetStoreLimit(i, storelimit.RemovePeer, 1)
+	}
+
+	// Add regions 1~4.
+	seq := newSequencer(3)
+	// Region 1 has the same distribution with the Region 2, which is used to test selectPeerToReplace.
+	tc.AddLeaderRegion(1, 1, 2, 3)
+	for i := uint64(2); i <= 5; i++ {
+		tc.AddLeaderRegion(i, seq.next(), seq.next(), seq.next())
+	}
+
+	scatterer := NewRegionScatterer(ctx, tc, oc)
+
+	for i := uint64(1); i <= 5; i++ {
+		region := tc.GetRegion(i)
+		if op, _ := scatterer.Scatter(region, ""); op != nil {
+			re.True(oc.AddOperator(op))
+		}
+	}
+}