From e48894eb7ff5fd1799386bb7a694d8faefb3c26e Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 22 Jun 2026 23:59:07 +0800 Subject: [PATCH 1/9] schedule: respect limits in concurrent patrol workers Signed-off-by: Ryan Leung --- pkg/schedule/checker/checker_controller.go | 218 ++++++++++++++---- .../checker/checker_controller_test.go | 64 +++++ server/cluster/cluster_test.go | 45 ++++ 3 files changed, 277 insertions(+), 50 deletions(-) create mode 100644 pkg/schedule/checker/checker_controller_test.go diff --git a/pkg/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index 7da5ae7bb49..accbdb2bd50 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -35,8 +35,10 @@ import ( sche "github.com/tikv/pd/pkg/schedule/core" "github.com/tikv/pd/pkg/schedule/labeler" "github.com/tikv/pd/pkg/schedule/operator" + "github.com/tikv/pd/pkg/schedule/types" "github.com/tikv/pd/pkg/utils/keyutil" "github.com/tikv/pd/pkg/utils/logutil" + "github.com/tikv/pd/pkg/utils/syncutil" ) const ( @@ -83,7 +85,8 @@ type Controller struct { // duration is the duration of the last patrol round. // It's exported, so it should be protected by a mutex. - duration atomic.Value // Store as time.Duration + duration atomic.Value // Store as time.Duration + operatorMu syncutil.Mutex // interval is the config interval of patrol regions. // It's used to update the ticker, so we need to @@ -118,6 +121,7 @@ func NewController(ctx context.Context, cluster sche.CheckerCluster, conf config suspectKeyRanges: cache.NewStringTTL(ctx, time.Minute, 3*time.Minute), patrolRegionContext: &PatrolRegionContext{}, interval: cluster.GetCheckerConfig().GetPatrolRegionInterval(), + workerCount: cluster.GetCheckerConfig().GetPatrolRegionWorkerCount(), patrolRegionScanLimit: calculateScanLimit(cluster), metrics: newCheckerControllerMetrics(), } @@ -147,20 +151,14 @@ func (c *Controller) PatrolRegions() { failpoint.Continue() }) c.updateTickerIfNeeded(ticker) - c.updatePatrolWorkersIfNeeded() - if c.cluster.IsSchedulingHalted() { - for len(c.patrolRegionContext.regionChan) > 0 { - <-c.patrolRegionContext.regionChan - } - log.Debug("skip patrol regions due to scheduling is halted") + // wait for the previous batch to be fully processed. + if !c.patrolRegionContext.isIdle() { continue } c.metrics.patrolRegionChannelSize.Set(float64(len(c.patrolRegionContext.regionChan))) - - // wait for the regionChan to be drained - waitDrainChanel := time.Now() - if len(c.patrolRegionContext.regionChan) > 0 { - c.metrics.patrolPhaseHistograms[phaseWaitForChannel].Observe(time.Since(waitDrainChanel).Seconds()) + c.updatePatrolWorkersIfNeeded() + if c.cluster.IsSchedulingHalted() { + log.Debug("skip patrol regions due to scheduling is halted") continue } @@ -189,6 +187,11 @@ func (c *Controller) PatrolRegions() { }) // When the key is nil, it means that the scan is finished. if len(key) == 0 { + if !c.patrolRegionContext.waitIdle(c.ctx) { + patrolCheckRegionsGauge.Set(0) + c.setPatrolRegionsDuration(0) + return + } // update the scan limit. c.patrolRegionScanLimit = calculateScanLimit(c.cluster) // update the metrics. @@ -224,6 +227,9 @@ func (c *Controller) updateTickerIfNeeded(ticker *time.Ticker) { func (c *Controller) updatePatrolWorkersIfNeeded() { newWorkersCount := c.cluster.GetCheckerConfig().GetPatrolRegionWorkerCount() if c.workerCount != newWorkersCount { + if !c.patrolRegionContext.isIdle() { + return + } oldWorkersCount := c.workerCount c.workerCount = newWorkersCount // Stop the old workers and start the new workers. @@ -237,6 +243,16 @@ func (c *Controller) updatePatrolWorkersIfNeeded() { } } +type operatorControllerLimit struct { + kind operator.OpKind + checkerType types.CheckerSchedulerType +} + +type checkRegionResult struct { + ops []*operator.Operator + limit *operatorControllerLimit +} + // GetPatrolRegionsDuration returns the duration of the last patrol region round. func (c *Controller) GetPatrolRegionsDuration() time.Duration { return c.duration.Load().(time.Duration) @@ -255,7 +271,7 @@ func (c *Controller) checkRegions(startKey []byte) (key []byte, regions []*core. } for _, region := range regions { - c.patrolRegionContext.regionChan <- region + c.patrolRegionContext.submit(region) key = region.GetEndKey() } return @@ -281,14 +297,13 @@ func (c *Controller) checkPriorityRegions() { removes = append(removes, id) continue } - ops := c.CheckRegion(region) + result := c.buildRegionCheckResult(region) + ops := result.ops // it should skip if region needs to merge if len(ops) == 0 || ops[0].HasRelatedMergeRegion() { continue } - if !c.opController.ExceedStoreLimit(ops...) { - c.opController.AddWaitingOperator(ops...) - } + c.addWaitingOperators(result) } for _, v := range removes { c.RemovePriorityRegions(v) @@ -298,6 +313,10 @@ func (c *Controller) checkPriorityRegions() { // CheckRegion will check the region and add a new operator if needed. // The function is exposed for test purpose. func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { + return c.buildRegionCheckResult(region).ops +} + +func (c *Controller) buildRegionCheckResult(region *core.RegionInfo) checkRegionResult { // If PD has restarted, it needs to check learners added before and promote them. // Don't check isRaftLearnerEnabled cause it maybe disable learner feature but there are still some learners to promote. opController := c.opController @@ -305,17 +324,18 @@ func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { if ops := measureChecker(c.metrics.checkRegionHistograms[jointStateChecker], func() []*operator.Operator { return []*operator.Operator{c.jointStateChecker.Check(region)} }); len(ops) > 0 { - return ops + return checkRegionResult{ops: ops} } if ops := measureChecker(c.metrics.checkRegionHistograms[splitChecker], func() []*operator.Operator { return []*operator.Operator{c.splitChecker.Check(region)} }); len(ops) > 0 { - return ops + return checkRegionResult{ops: ops} } if c.conf.IsPlacementRulesEnabled() { - if ops := measureChecker(c.metrics.checkRegionHistograms[ruleChecker], func() []*operator.Operator { + var result checkRegionResult + measure(c.metrics.checkRegionHistograms[ruleChecker], func() { skipRuleCheck := c.cluster.GetCheckerConfig().IsPlacementRulesCacheEnabled() && c.cluster.GetRuleManager().IsRegionFitCached(c.cluster, region) if skipRuleCheck { @@ -330,33 +350,50 @@ func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { }) fit := c.priorityInspector.Inspect(region) if opController.OperatorCount(operator.OpReplica) < c.conf.GetReplicaScheduleLimit() { - return []*operator.Operator{c.ruleChecker.CheckWithFit(region, fit)} + if op := c.ruleChecker.CheckWithFit(region, fit); op != nil { + result = checkRegionResult{ + ops: []*operator.Operator{op}, + limit: &operatorControllerLimit{ + kind: operator.OpReplica, + checkerType: c.ruleChecker.GetType(), + }, + } + } + return } operator.IncOperatorLimitCounter(c.ruleChecker.GetType(), operator.OpReplica) c.pendingProcessedRegions.Put(region.GetID(), nil) } - return nil - }); len(ops) > 0 { - return ops + }) + if len(result.ops) > 0 { + return result } } else { if ops := measureChecker(c.metrics.checkRegionHistograms[learnerChecker], func() []*operator.Operator { return []*operator.Operator{c.learnerChecker.Check(region)} }); len(ops) > 0 { - return ops + return checkRegionResult{ops: ops} } - if ops := measureChecker(c.metrics.checkRegionHistograms[replicaChecker], func() []*operator.Operator { + var result checkRegionResult + measure(c.metrics.checkRegionHistograms[replicaChecker], func() { if op := c.replicaChecker.Check(region); op != nil { if opController.OperatorCount(operator.OpReplica) < c.conf.GetReplicaScheduleLimit() { - return []*operator.Operator{op} + result = checkRegionResult{ + ops: []*operator.Operator{op}, + limit: &operatorControllerLimit{ + kind: operator.OpReplica, + checkerType: c.replicaChecker.GetType(), + }, + } + return } operator.IncOperatorLimitCounter(c.replicaChecker.GetType(), operator.OpReplica) c.pendingProcessedRegions.Put(region.GetID(), nil) } - return nil - }); len(ops) > 0 { - return ops + }) + if len(result.ops) > 0 { + return result } } // skip the joint checker, split checker and rule checker when region label is set to "schedule=deny". @@ -364,33 +401,52 @@ func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { l := c.cluster.GetRegionLabeler() if l.ScheduleDisabled(region) { denyCheckersByLabelerCounter.Inc() - return nil + return checkRegionResult{} } - if ops := measureChecker(c.metrics.checkRegionHistograms[affinityChecker], func() []*operator.Operator { + var affinityResult checkRegionResult + measure(c.metrics.checkRegionHistograms[affinityChecker], func() { if opController.OperatorCount(operator.OpAffinity) < c.conf.GetAffinityScheduleLimit() { - // It makes sure that two affinity merge operators can be added successfully altogether. - return c.affinityChecker.Check(region) + if ops := c.affinityChecker.Check(region); len(ops) > 0 { + affinityResult = checkRegionResult{ + ops: ops, + limit: &operatorControllerLimit{ + kind: operator.OpAffinity, + checkerType: c.affinityChecker.GetType(), + }, + } + } + return } if c.affinityChecker.hasAffinityGroups() { operator.IncOperatorLimitCounter(c.affinityChecker.GetType(), operator.OpAffinity) } - return nil - }); len(ops) > 0 { - return ops + }) + if len(affinityResult.ops) > 0 { + return affinityResult } - if ops := measureChecker(c.metrics.checkRegionHistograms[mergeChecker], func() []*operator.Operator { + var mergeResult checkRegionResult + measure(c.metrics.checkRegionHistograms[mergeChecker], func() { if opController.OperatorCount(operator.OpMerge) < c.conf.GetMergeScheduleLimit() { - // It makes sure that two operators can be added successfully altogether. - return c.mergeChecker.Check(region) + if ops := c.mergeChecker.Check(region); len(ops) > 0 { + // It makes sure that two operators can be added successfully altogether. + mergeResult = checkRegionResult{ + ops: ops, + limit: &operatorControllerLimit{ + kind: operator.OpMerge, + checkerType: c.mergeChecker.GetType(), + }, + } + } + return } operator.IncOperatorLimitCounter(c.mergeChecker.GetType(), operator.OpMerge) - return nil - }); len(ops) > 0 { - return ops + }) + if len(mergeResult.ops) > 0 { + return mergeResult } - return nil + return checkRegionResult{} } func (c *Controller) tryAddOperators(region *core.RegionInfo) { @@ -398,24 +454,61 @@ func (c *Controller) tryAddOperators(region *core.RegionInfo) { // the region could be recent split, continue to wait. return } + if c.cluster.IsSchedulingHalted() { + return + } id := region.GetID() if c.opController.GetOperator(id) != nil { c.RemovePendingProcessedRegion(id) return } - ops := c.CheckRegion(region) - if len(ops) == 0 { + result := c.buildRegionCheckResult(region) + if len(result.ops) == 0 { return } - if !c.opController.ExceedStoreLimit(ops...) { - c.opController.AddWaitingOperator(ops...) + if c.addWaitingOperators(result) { c.RemovePendingProcessedRegion(id) } else { + if c.cluster.IsSchedulingHalted() { + return + } c.AddPendingProcessedRegions(true, id) } } +func (c *Controller) addWaitingOperators(result checkRegionResult) bool { + if len(result.ops) == 0 { + return false + } + c.operatorMu.Lock() + defer c.operatorMu.Unlock() + if c.cluster.IsSchedulingHalted() { + return false + } + if result.limit != nil && c.opController.OperatorCount(result.limit.kind) >= c.getOperatorLimit(result.limit.kind) { + operator.IncOperatorLimitCounter(result.limit.checkerType, result.limit.kind) + return false + } + if c.opController.ExceedStoreLimit(result.ops...) { + return false + } + return c.opController.AddWaitingOperator(result.ops...) > 0 +} + +func (c *Controller) getOperatorLimit(kind operator.OpKind) uint64 { + switch kind { + case operator.OpReplica: + return c.conf.GetReplicaScheduleLimit() + case operator.OpMerge: + return c.conf.GetMergeScheduleLimit() + case operator.OpAffinity: + return c.conf.GetAffinityScheduleLimit() + default: + return 0 + } +} + // GetMergeChecker returns the merge checker. func (c *Controller) GetMergeChecker() *MergeChecker { return c.mergeChecker @@ -566,7 +659,7 @@ func (c *Controller) IsPatrolRegionChanEmpty() bool { if c.patrolRegionContext == nil { return true } - return len(c.patrolRegionContext.regionChan) == 0 + return c.patrolRegionContext.isIdle() } // PatrolRegionContext is used to store the context of patrol regions. @@ -575,11 +668,13 @@ type PatrolRegionContext struct { workersCancel context.CancelFunc regionChan chan *core.RegionInfo wg sync.WaitGroup + pendingCount atomic.Int64 } func (p *PatrolRegionContext) init(ctx context.Context) { p.regionChan = make(chan *core.RegionInfo, patrolRegionChanLen) p.workersCtx, p.workersCancel = context.WithCancel(ctx) + p.pendingCount.Store(0) } func (p *PatrolRegionContext) stop() { @@ -603,7 +698,10 @@ func (p *PatrolRegionContext) startPatrolRegionWorkers(c *Controller) { log.Debug("region channel is closed", zap.Int("worker-id", i)) return } - c.tryAddOperators(region) + func() { + defer p.pendingCount.Add(-1) + c.tryAddOperators(region) + }() case <-p.workersCtx.Done(): log.Debug("region worker is closed", zap.Int("worker-id", i)) return @@ -613,6 +711,26 @@ func (p *PatrolRegionContext) startPatrolRegionWorkers(c *Controller) { } } +func (p *PatrolRegionContext) submit(region *core.RegionInfo) { + p.pendingCount.Add(1) + p.regionChan <- region +} + +func (p *PatrolRegionContext) isIdle() bool { + return p.pendingCount.Load() == 0 +} + +func (p *PatrolRegionContext) waitIdle(ctx context.Context) bool { + for !p.isIdle() { + select { + case <-ctx.Done(): + return false + case <-time.After(time.Millisecond): + } + } + return true +} + // GetPatrolRegionScanLimit returns the limit of regions to scan. // It only used for test. func (c *Controller) GetPatrolRegionScanLimit() int { diff --git a/pkg/schedule/checker/checker_controller_test.go b/pkg/schedule/checker/checker_controller_test.go new file mode 100644 index 00000000000..025e0e1a009 --- /dev/null +++ b/pkg/schedule/checker/checker_controller_test.go @@ -0,0 +1,64 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package checker + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/tikv/pd/pkg/mock/mockcluster" + "github.com/tikv/pd/pkg/mock/mockconfig" + "github.com/tikv/pd/pkg/schedule/hbstream" + "github.com/tikv/pd/pkg/schedule/operator" + "github.com/tikv/pd/pkg/schedule/types" +) + +func TestAddWaitingOperatorsSkipsSchedulingHalted(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + opt := mockconfig.NewTestOptions() + tc := mockcluster.NewCluster(ctx, opt) + for storeID := uint64(1); storeID <= 3; storeID++ { + tc.AddRegionStore(storeID, 0) + } + region := tc.AddLeaderRegion(1, 1, 2) + + stream := hbstream.NewTestHeartbeatStreams(ctx, tc, false) + defer stream.Close() + oc := operator.NewController(ctx, tc.GetBasicCluster(), tc.GetSharedConfig(), stream) + controller := NewController(ctx, tc, tc.GetCheckerConfig(), oc) + result := checkRegionResult{ + ops: []*operator.Operator{ + operator.NewTestOperator(region.GetID(), region.GetRegionEpoch(), operator.OpReplica), + }, + limit: &operatorControllerLimit{ + kind: operator.OpReplica, + checkerType: types.ReplicaChecker, + }, + } + + tc.SetHaltScheduling(true, "test") + re.False(controller.addWaitingOperators(result)) + re.Empty(oc.GetOperators()) + re.Empty(oc.GetWaitingOperators()) + + tc.SetHaltScheduling(false, "test") + re.True(controller.addWaitingOperators(result)) + re.Len(oc.GetOperators(), 1) +} diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go index b8080d7a45f..f492118cee9 100644 --- a/server/cluster/cluster_test.go +++ b/server/cluster/cluster_test.go @@ -3195,6 +3195,51 @@ func TestPatrolRegionConcurrency(t *testing.T) { re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/schedule/checker/breakPatrol")) } +func TestPatrolRegionConcurrencyRespectReplicaLimit(t *testing.T) { + re := require.New(t) + + const ( + regionNum = 128 + replicaScheduleLimit = 1 + ) + + tc, co, cleanup := prepare(func(cfg *sc.ScheduleConfig) { + cfg.PatrolRegionWorkerCount = 8 + cfg.ReplicaScheduleLimit = replicaScheduleLimit + }, nil, nil, re) + oc := co.GetOperatorController() + checker := co.GetCheckerController() + done := make(chan struct{}) + defer func() { + cleanup() + select { + case <-done: + case <-time.After(time.Second): + re.Fail("patrol regions did not stop") + } + }() + + for i := range 3 { + re.NoError(tc.addRegionStore(uint64(i+1), regionNum)) + } + for i := range regionNum { + re.NoError(tc.addLeaderRegion(uint64(i+1), 1)) + } + + go func() { + checker.PatrolRegions() + close(done) + }() + + testutil.Eventually(re, func() bool { + count := oc.OperatorCount(operator.OpReplica) + return count > uint64(replicaScheduleLimit) || + (count == uint64(replicaScheduleLimit) && len(checker.GetPendingProcessedRegions()) > 0) + }) + re.LessOrEqual(oc.OperatorCount(operator.OpReplica), uint64(replicaScheduleLimit)) + re.Len(oc.GetOperators(), replicaScheduleLimit) +} + func checkOperatorDuplicate(re *require.Assertions, ops []*operator.Operator) { regionMap := make(map[uint64]struct{}) for _, op := range ops { From 16c81444438662be2cada099e4ebc132a4312239 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 17:14:38 +0800 Subject: [PATCH 2/9] schedule: clean up patrol operator admission Signed-off-by: Ryan Leung --- pkg/schedule/checker/checker_controller.go | 197 +++++++++--------- .../checker/checker_controller_test.go | 17 +- 2 files changed, 109 insertions(+), 105 deletions(-) diff --git a/pkg/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index accbdb2bd50..3401bd675f9 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -243,14 +243,31 @@ func (c *Controller) updatePatrolWorkersIfNeeded() { } } -type operatorControllerLimit struct { +// operatorLimit is rechecked while admitting a candidate into the operator queue. +type operatorLimit struct { kind operator.OpKind checkerType types.CheckerSchedulerType + getLimit func() uint64 } -type checkRegionResult struct { +func (l operatorLimit) reached(opController *operator.Controller) bool { + return l.getLimit != nil && opController.OperatorCount(l.kind) >= l.getLimit() +} + +func (l operatorLimit) recordLimit() { + if l.getLimit != nil { + operator.IncOperatorLimitCounter(l.checkerType, l.kind) + } +} + +// operatorCandidate is a checker result that still needs operator-controller admission. +type operatorCandidate struct { ops []*operator.Operator - limit *operatorControllerLimit + limit operatorLimit +} + +func newOperatorCandidate(limit operatorLimit, ops ...*operator.Operator) operatorCandidate { + return operatorCandidate{ops: ops, limit: limit} } // GetPatrolRegionsDuration returns the duration of the last patrol region round. @@ -297,13 +314,13 @@ func (c *Controller) checkPriorityRegions() { removes = append(removes, id) continue } - result := c.buildRegionCheckResult(region) - ops := result.ops + candidate := c.buildOperatorCandidate(region) + ops := candidate.ops // it should skip if region needs to merge if len(ops) == 0 || ops[0].HasRelatedMergeRegion() { continue } - c.addWaitingOperators(result) + c.admitWaitingOperators(candidate) } for _, v := range removes { c.RemovePriorityRegions(v) @@ -313,28 +330,55 @@ func (c *Controller) checkPriorityRegions() { // CheckRegion will check the region and add a new operator if needed. // The function is exposed for test purpose. func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { - return c.buildRegionCheckResult(region).ops + return c.buildOperatorCandidate(region).ops +} + +func (c *Controller) replicaLimit(checkerType types.CheckerSchedulerType) operatorLimit { + return operatorLimit{ + kind: operator.OpReplica, + checkerType: checkerType, + getLimit: c.conf.GetReplicaScheduleLimit, + } +} + +func (c *Controller) affinityLimit() operatorLimit { + return operatorLimit{ + kind: operator.OpAffinity, + checkerType: c.affinityChecker.GetType(), + getLimit: c.conf.GetAffinityScheduleLimit, + } +} + +func (c *Controller) mergeLimit() operatorLimit { + return operatorLimit{ + kind: operator.OpMerge, + checkerType: c.mergeChecker.GetType(), + getLimit: c.conf.GetMergeScheduleLimit, + } +} + +func (c *Controller) deferByLimit(regionID uint64, limit operatorLimit) { + limit.recordLimit() + c.pendingProcessedRegions.Put(regionID, nil) } -func (c *Controller) buildRegionCheckResult(region *core.RegionInfo) checkRegionResult { +func (c *Controller) buildOperatorCandidate(region *core.RegionInfo) operatorCandidate { // If PD has restarted, it needs to check learners added before and promote them. // Don't check isRaftLearnerEnabled cause it maybe disable learner feature but there are still some learners to promote. - opController := c.opController - if ops := measureChecker(c.metrics.checkRegionHistograms[jointStateChecker], func() []*operator.Operator { return []*operator.Operator{c.jointStateChecker.Check(region)} }); len(ops) > 0 { - return checkRegionResult{ops: ops} + return newOperatorCandidate(operatorLimit{}, ops...) } if ops := measureChecker(c.metrics.checkRegionHistograms[splitChecker], func() []*operator.Operator { return []*operator.Operator{c.splitChecker.Check(region)} }); len(ops) > 0 { - return checkRegionResult{ops: ops} + return newOperatorCandidate(operatorLimit{}, ops...) } if c.conf.IsPlacementRulesEnabled() { - var result checkRegionResult + var candidate operatorCandidate measure(c.metrics.checkRegionHistograms[ruleChecker], func() { skipRuleCheck := c.cluster.GetCheckerConfig().IsPlacementRulesCacheEnabled() && c.cluster.GetRuleManager().IsRegionFitCached(c.cluster, region) @@ -349,51 +393,39 @@ func (c *Controller) buildRegionCheckResult(region *core.RegionInfo) checkRegion panic("cached should be used") }) fit := c.priorityInspector.Inspect(region) - if opController.OperatorCount(operator.OpReplica) < c.conf.GetReplicaScheduleLimit() { - if op := c.ruleChecker.CheckWithFit(region, fit); op != nil { - result = checkRegionResult{ - ops: []*operator.Operator{op}, - limit: &operatorControllerLimit{ - kind: operator.OpReplica, - checkerType: c.ruleChecker.GetType(), - }, - } - } + limit := c.replicaLimit(c.ruleChecker.GetType()) + if limit.reached(c.opController) { + c.deferByLimit(region.GetID(), limit) return } - operator.IncOperatorLimitCounter(c.ruleChecker.GetType(), operator.OpReplica) - c.pendingProcessedRegions.Put(region.GetID(), nil) + if op := c.ruleChecker.CheckWithFit(region, fit); op != nil { + candidate = newOperatorCandidate(limit, op) + } } }) - if len(result.ops) > 0 { - return result + if len(candidate.ops) > 0 { + return candidate } } else { if ops := measureChecker(c.metrics.checkRegionHistograms[learnerChecker], func() []*operator.Operator { return []*operator.Operator{c.learnerChecker.Check(region)} }); len(ops) > 0 { - return checkRegionResult{ops: ops} + return newOperatorCandidate(operatorLimit{}, ops...) } - var result checkRegionResult + var candidate operatorCandidate measure(c.metrics.checkRegionHistograms[replicaChecker], func() { if op := c.replicaChecker.Check(region); op != nil { - if opController.OperatorCount(operator.OpReplica) < c.conf.GetReplicaScheduleLimit() { - result = checkRegionResult{ - ops: []*operator.Operator{op}, - limit: &operatorControllerLimit{ - kind: operator.OpReplica, - checkerType: c.replicaChecker.GetType(), - }, - } + limit := c.replicaLimit(c.replicaChecker.GetType()) + if limit.reached(c.opController) { + c.deferByLimit(region.GetID(), limit) return } - operator.IncOperatorLimitCounter(c.replicaChecker.GetType(), operator.OpReplica) - c.pendingProcessedRegions.Put(region.GetID(), nil) + candidate = newOperatorCandidate(limit, op) } }) - if len(result.ops) > 0 { - return result + if len(candidate.ops) > 0 { + return candidate } } // skip the joint checker, split checker and rule checker when region label is set to "schedule=deny". @@ -401,52 +433,42 @@ func (c *Controller) buildRegionCheckResult(region *core.RegionInfo) checkRegion l := c.cluster.GetRegionLabeler() if l.ScheduleDisabled(region) { denyCheckersByLabelerCounter.Inc() - return checkRegionResult{} + return operatorCandidate{} } - var affinityResult checkRegionResult + affinityLimit := c.affinityLimit() + var affinityCandidate operatorCandidate measure(c.metrics.checkRegionHistograms[affinityChecker], func() { - if opController.OperatorCount(operator.OpAffinity) < c.conf.GetAffinityScheduleLimit() { - if ops := c.affinityChecker.Check(region); len(ops) > 0 { - affinityResult = checkRegionResult{ - ops: ops, - limit: &operatorControllerLimit{ - kind: operator.OpAffinity, - checkerType: c.affinityChecker.GetType(), - }, - } + if affinityLimit.reached(c.opController) { + if c.affinityChecker.hasAffinityGroups() { + affinityLimit.recordLimit() } return } - if c.affinityChecker.hasAffinityGroups() { - operator.IncOperatorLimitCounter(c.affinityChecker.GetType(), operator.OpAffinity) + if ops := c.affinityChecker.Check(region); len(ops) > 0 { + affinityCandidate = newOperatorCandidate(affinityLimit, ops...) } }) - if len(affinityResult.ops) > 0 { - return affinityResult + if len(affinityCandidate.ops) > 0 { + return affinityCandidate } - var mergeResult checkRegionResult + mergeLimit := c.mergeLimit() + var mergeCandidate operatorCandidate measure(c.metrics.checkRegionHistograms[mergeChecker], func() { - if opController.OperatorCount(operator.OpMerge) < c.conf.GetMergeScheduleLimit() { - if ops := c.mergeChecker.Check(region); len(ops) > 0 { - // It makes sure that two operators can be added successfully altogether. - mergeResult = checkRegionResult{ - ops: ops, - limit: &operatorControllerLimit{ - kind: operator.OpMerge, - checkerType: c.mergeChecker.GetType(), - }, - } - } + if mergeLimit.reached(c.opController) { + mergeLimit.recordLimit() return } - operator.IncOperatorLimitCounter(c.mergeChecker.GetType(), operator.OpMerge) + if ops := c.mergeChecker.Check(region); len(ops) > 0 { + // It makes sure that two operators can be added successfully altogether. + mergeCandidate = newOperatorCandidate(mergeLimit, ops...) + } }) - if len(mergeResult.ops) > 0 { - return mergeResult + if len(mergeCandidate.ops) > 0 { + return mergeCandidate } - return checkRegionResult{} + return operatorCandidate{} } func (c *Controller) tryAddOperators(region *core.RegionInfo) { @@ -462,12 +484,12 @@ func (c *Controller) tryAddOperators(region *core.RegionInfo) { c.RemovePendingProcessedRegion(id) return } - result := c.buildRegionCheckResult(region) - if len(result.ops) == 0 { + candidate := c.buildOperatorCandidate(region) + if len(candidate.ops) == 0 { return } - if c.addWaitingOperators(result) { + if c.admitWaitingOperators(candidate) { c.RemovePendingProcessedRegion(id) } else { if c.cluster.IsSchedulingHalted() { @@ -477,8 +499,8 @@ func (c *Controller) tryAddOperators(region *core.RegionInfo) { } } -func (c *Controller) addWaitingOperators(result checkRegionResult) bool { - if len(result.ops) == 0 { +func (c *Controller) admitWaitingOperators(candidate operatorCandidate) bool { + if len(candidate.ops) == 0 { return false } c.operatorMu.Lock() @@ -486,27 +508,14 @@ func (c *Controller) addWaitingOperators(result checkRegionResult) bool { if c.cluster.IsSchedulingHalted() { return false } - if result.limit != nil && c.opController.OperatorCount(result.limit.kind) >= c.getOperatorLimit(result.limit.kind) { - operator.IncOperatorLimitCounter(result.limit.checkerType, result.limit.kind) + if candidate.limit.reached(c.opController) { + candidate.limit.recordLimit() return false } - if c.opController.ExceedStoreLimit(result.ops...) { + if c.opController.ExceedStoreLimit(candidate.ops...) { return false } - return c.opController.AddWaitingOperator(result.ops...) > 0 -} - -func (c *Controller) getOperatorLimit(kind operator.OpKind) uint64 { - switch kind { - case operator.OpReplica: - return c.conf.GetReplicaScheduleLimit() - case operator.OpMerge: - return c.conf.GetMergeScheduleLimit() - case operator.OpAffinity: - return c.conf.GetAffinityScheduleLimit() - default: - return 0 - } + return c.opController.AddWaitingOperator(candidate.ops...) > 0 } // GetMergeChecker returns the merge checker. diff --git a/pkg/schedule/checker/checker_controller_test.go b/pkg/schedule/checker/checker_controller_test.go index 025e0e1a009..2851bcd60fe 100644 --- a/pkg/schedule/checker/checker_controller_test.go +++ b/pkg/schedule/checker/checker_controller_test.go @@ -43,22 +43,17 @@ func TestAddWaitingOperatorsSkipsSchedulingHalted(t *testing.T) { defer stream.Close() oc := operator.NewController(ctx, tc.GetBasicCluster(), tc.GetSharedConfig(), stream) controller := NewController(ctx, tc, tc.GetCheckerConfig(), oc) - result := checkRegionResult{ - ops: []*operator.Operator{ - operator.NewTestOperator(region.GetID(), region.GetRegionEpoch(), operator.OpReplica), - }, - limit: &operatorControllerLimit{ - kind: operator.OpReplica, - checkerType: types.ReplicaChecker, - }, - } + candidate := newOperatorCandidate( + controller.replicaLimit(types.ReplicaChecker), + operator.NewTestOperator(region.GetID(), region.GetRegionEpoch(), operator.OpReplica), + ) tc.SetHaltScheduling(true, "test") - re.False(controller.addWaitingOperators(result)) + re.False(controller.admitWaitingOperators(candidate)) re.Empty(oc.GetOperators()) re.Empty(oc.GetWaitingOperators()) tc.SetHaltScheduling(false, "test") - re.True(controller.addWaitingOperators(result)) + re.True(controller.admitWaitingOperators(candidate)) re.Len(oc.GetOperators(), 1) } From e946a88b14494ea17372748a8a570b2e0f990c9e Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 17:52:09 +0800 Subject: [PATCH 3/9] tests: wait for affinity group visibility Signed-off-by: Ryan Leung --- tests/server/apiv2/handlers/affinity_test.go | 29 ++++++++++++++------ 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/tests/server/apiv2/handlers/affinity_test.go b/tests/server/apiv2/handlers/affinity_test.go index 755d1362247..697d0d33565 100644 --- a/tests/server/apiv2/handlers/affinity_test.go +++ b/tests/server/apiv2/handlers/affinity_test.go @@ -139,6 +139,24 @@ func mustGetAllAffinityGroups(re *require.Assertions, serverAddr string) *handle return &result } +func waitForAffinityGroups(re *require.Assertions, url string, groupIDs ...string) *handlers.AffinityGroupsResponse { + var result handlers.AffinityGroupsResponse + testutil.Eventually(re, func() bool { + result = handlers.AffinityGroupsResponse{} + err := testutil.ReadGetJSON(re, tests.TestDialClient, url, &result) + if err != nil || len(result.AffinityGroups) != len(groupIDs) { + return false + } + for _, groupID := range groupIDs { + if _, ok := result.AffinityGroups[groupID]; !ok { + return false + } + } + return true + }) + return &result +} + func mustPutHealthyStore(re *require.Assertions, cluster *tests.TestCluster, store *metapb.Store) { tests.MustPutStore(re, cluster, store) @@ -660,9 +678,7 @@ func (suite *affinityHandlerTestSuite) TestAffinityGroupCreateSkipExistCheck() { re.Contains(result.AffinityGroups, "existing") re.Contains(result.AffinityGroups, "new") - listResp := mustGetAllAffinityGroups(re, serverAddr) - re.Contains(listResp.AffinityGroups, "existing") - re.Contains(listResp.AffinityGroups, "new") + waitForAffinityGroups(re, getAffinityGroupURL(serverAddr), "existing", "new") }) } @@ -855,12 +871,7 @@ func (suite *affinityHandlerTestSuite) TestAffinityListWithIDs() { } mustCreateAffinityGroups(re, serverAddr, &createReq) - var result handlers.AffinityGroupsResponse - err := testutil.ReadGetJSON(re, tests.TestDialClient, getAffinityGroupURL(serverAddr)+"?ids=group-1&ids=group-3&ids=missing", &result) - re.NoError(err) - re.Len(result.AffinityGroups, 2) - re.Contains(result.AffinityGroups, "group-1") - re.Contains(result.AffinityGroups, "group-3") + waitForAffinityGroups(re, getAffinityGroupURL(serverAddr)+"?ids=group-1&ids=group-3&ids=missing", "group-1", "group-3") }) } From 71fe19b9debf3671a025195863bf14e884e8d535 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 18:21:35 +0800 Subject: [PATCH 4/9] tests: avoid unused affinity wait result Signed-off-by: Ryan Leung --- tests/server/apiv2/handlers/affinity_test.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/server/apiv2/handlers/affinity_test.go b/tests/server/apiv2/handlers/affinity_test.go index 697d0d33565..f84eb53b8ed 100644 --- a/tests/server/apiv2/handlers/affinity_test.go +++ b/tests/server/apiv2/handlers/affinity_test.go @@ -139,7 +139,7 @@ func mustGetAllAffinityGroups(re *require.Assertions, serverAddr string) *handle return &result } -func waitForAffinityGroups(re *require.Assertions, url string, groupIDs ...string) *handlers.AffinityGroupsResponse { +func waitForAffinityGroups(re *require.Assertions, url string, groupIDs ...string) { var result handlers.AffinityGroupsResponse testutil.Eventually(re, func() bool { result = handlers.AffinityGroupsResponse{} @@ -154,7 +154,6 @@ func waitForAffinityGroups(re *require.Assertions, url string, groupIDs ...strin } return true }) - return &result } func mustPutHealthyStore(re *require.Assertions, cluster *tests.TestCluster, store *metapb.Store) { From c8040faa52d397eea185736556c9c34570a9fb89 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 18:51:31 +0800 Subject: [PATCH 5/9] tests: avoid assertions inside affinity wait loop Signed-off-by: Ryan Leung --- tests/server/apiv2/handlers/affinity_test.go | 23 ++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/tests/server/apiv2/handlers/affinity_test.go b/tests/server/apiv2/handlers/affinity_test.go index f84eb53b8ed..7e7784137a1 100644 --- a/tests/server/apiv2/handlers/affinity_test.go +++ b/tests/server/apiv2/handlers/affinity_test.go @@ -141,10 +141,24 @@ func mustGetAllAffinityGroups(re *require.Assertions, serverAddr string) *handle func waitForAffinityGroups(re *require.Assertions, url string, groupIDs ...string) { var result handlers.AffinityGroupsResponse + var lastBody []byte + var lastErr error + var lastStatusCode int testutil.Eventually(re, func() bool { result = handlers.AffinityGroupsResponse{} - err := testutil.ReadGetJSON(re, tests.TestDialClient, url, &result) - if err != nil || len(result.AffinityGroups) != len(groupIDs) { + resp, err := apiutil.GetJSON(tests.TestDialClient, url, nil) + if err != nil { + lastErr = err + return false + } + defer resp.Body.Close() + lastStatusCode = resp.StatusCode + lastBody, lastErr = io.ReadAll(resp.Body) + if lastErr != nil || resp.StatusCode != http.StatusOK { + return false + } + lastErr = json.Unmarshal(lastBody, &result) + if lastErr != nil || len(result.AffinityGroups) != len(groupIDs) { return false } for _, groupID := range groupIDs { @@ -154,6 +168,11 @@ func waitForAffinityGroups(re *require.Assertions, url string, groupIDs ...strin } return true }) + re.NoError(lastErr) + re.Equal(http.StatusOK, lastStatusCode, "resp: "+string(lastBody)) + for _, groupID := range groupIDs { + re.Contains(result.AffinityGroups, groupID) + } } func mustPutHealthyStore(re *require.Assertions, cluster *tests.TestCluster, store *metapb.Store) { From 7118142c7176fa718319e18451dec5ed6592f6e8 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Thu, 9 Jul 2026 15:00:12 +0800 Subject: [PATCH 6/9] schedule: report patrol pending work Signed-off-by: Ryan Leung --- pkg/schedule/checker/checker_controller.go | 17 +++++-- tests/server/apiv2/handlers/affinity_test.go | 47 ++++---------------- 2 files changed, 22 insertions(+), 42 deletions(-) diff --git a/pkg/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index 3401bd675f9..66b94cd9d70 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -151,11 +151,12 @@ func (c *Controller) PatrolRegions() { failpoint.Continue() }) c.updateTickerIfNeeded(ticker) + pendingRegionCount := c.patrolRegionContext.pendingRegionCount() + c.metrics.patrolRegionChannelSize.Set(float64(pendingRegionCount)) // wait for the previous batch to be fully processed. - if !c.patrolRegionContext.isIdle() { + if pendingRegionCount > 0 { continue } - c.metrics.patrolRegionChannelSize.Set(float64(len(c.patrolRegionContext.regionChan))) c.updatePatrolWorkersIfNeeded() if c.cluster.IsSchedulingHalted() { log.Debug("skip patrol regions due to scheduling is halted") @@ -187,7 +188,11 @@ func (c *Controller) PatrolRegions() { }) // When the key is nil, it means that the scan is finished. if len(key) == 0 { - if !c.patrolRegionContext.waitIdle(c.ctx) { + var idle bool + measure(c.metrics.patrolPhaseHistograms[phaseWaitForChannel], func() { + idle = c.patrolRegionContext.waitIdle(c.ctx) + }) + if !idle { patrolCheckRegionsGauge.Set(0) c.setPatrolRegionsDuration(0) return @@ -725,8 +730,12 @@ func (p *PatrolRegionContext) submit(region *core.RegionInfo) { p.regionChan <- region } +func (p *PatrolRegionContext) pendingRegionCount() int64 { + return p.pendingCount.Load() +} + func (p *PatrolRegionContext) isIdle() bool { - return p.pendingCount.Load() == 0 + return p.pendingRegionCount() == 0 } func (p *PatrolRegionContext) waitIdle(ctx context.Context) bool { diff --git a/tests/server/apiv2/handlers/affinity_test.go b/tests/server/apiv2/handlers/affinity_test.go index 7e7784137a1..755d1362247 100644 --- a/tests/server/apiv2/handlers/affinity_test.go +++ b/tests/server/apiv2/handlers/affinity_test.go @@ -139,42 +139,6 @@ func mustGetAllAffinityGroups(re *require.Assertions, serverAddr string) *handle return &result } -func waitForAffinityGroups(re *require.Assertions, url string, groupIDs ...string) { - var result handlers.AffinityGroupsResponse - var lastBody []byte - var lastErr error - var lastStatusCode int - testutil.Eventually(re, func() bool { - result = handlers.AffinityGroupsResponse{} - resp, err := apiutil.GetJSON(tests.TestDialClient, url, nil) - if err != nil { - lastErr = err - return false - } - defer resp.Body.Close() - lastStatusCode = resp.StatusCode - lastBody, lastErr = io.ReadAll(resp.Body) - if lastErr != nil || resp.StatusCode != http.StatusOK { - return false - } - lastErr = json.Unmarshal(lastBody, &result) - if lastErr != nil || len(result.AffinityGroups) != len(groupIDs) { - return false - } - for _, groupID := range groupIDs { - if _, ok := result.AffinityGroups[groupID]; !ok { - return false - } - } - return true - }) - re.NoError(lastErr) - re.Equal(http.StatusOK, lastStatusCode, "resp: "+string(lastBody)) - for _, groupID := range groupIDs { - re.Contains(result.AffinityGroups, groupID) - } -} - func mustPutHealthyStore(re *require.Assertions, cluster *tests.TestCluster, store *metapb.Store) { tests.MustPutStore(re, cluster, store) @@ -696,7 +660,9 @@ func (suite *affinityHandlerTestSuite) TestAffinityGroupCreateSkipExistCheck() { re.Contains(result.AffinityGroups, "existing") re.Contains(result.AffinityGroups, "new") - waitForAffinityGroups(re, getAffinityGroupURL(serverAddr), "existing", "new") + listResp := mustGetAllAffinityGroups(re, serverAddr) + re.Contains(listResp.AffinityGroups, "existing") + re.Contains(listResp.AffinityGroups, "new") }) } @@ -889,7 +855,12 @@ func (suite *affinityHandlerTestSuite) TestAffinityListWithIDs() { } mustCreateAffinityGroups(re, serverAddr, &createReq) - waitForAffinityGroups(re, getAffinityGroupURL(serverAddr)+"?ids=group-1&ids=group-3&ids=missing", "group-1", "group-3") + var result handlers.AffinityGroupsResponse + err := testutil.ReadGetJSON(re, tests.TestDialClient, getAffinityGroupURL(serverAddr)+"?ids=group-1&ids=group-3&ids=missing", &result) + re.NoError(err) + re.Len(result.AffinityGroups, 2) + re.Contains(result.AffinityGroups, "group-1") + re.Contains(result.AffinityGroups, "group-3") }) } From 919f8c808963a1402a5bd37c26b90487cbc90d92 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Thu, 9 Jul 2026 17:13:17 +0800 Subject: [PATCH 7/9] schedule: clarify patrol operator checks Signed-off-by: Ryan Leung --- pkg/schedule/checker/checker_controller.go | 204 ++++++++++++--------- pkg/schedule/checker/metrics.go | 12 +- 2 files changed, 127 insertions(+), 89 deletions(-) diff --git a/pkg/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index 66b94cd9d70..4648b7608bb 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -151,10 +151,7 @@ func (c *Controller) PatrolRegions() { failpoint.Continue() }) c.updateTickerIfNeeded(ticker) - pendingRegionCount := c.patrolRegionContext.pendingRegionCount() - c.metrics.patrolRegionChannelSize.Set(float64(pendingRegionCount)) - // wait for the previous batch to be fully processed. - if pendingRegionCount > 0 { + if c.recordPendingPatrolWork() > 0 { continue } c.updatePatrolWorkersIfNeeded() @@ -188,22 +185,11 @@ func (c *Controller) PatrolRegions() { }) // When the key is nil, it means that the scan is finished. if len(key) == 0 { - var idle bool - measure(c.metrics.patrolPhaseHistograms[phaseWaitForChannel], func() { - idle = c.patrolRegionContext.waitIdle(c.ctx) - }) - if !idle { - patrolCheckRegionsGauge.Set(0) - c.setPatrolRegionsDuration(0) + var completed bool + start, completed = c.finishPatrolRound(start) + if !completed { return } - // update the scan limit. - c.patrolRegionScanLimit = calculateScanLimit(c.cluster) - // update the metrics. - dur := time.Since(start) - patrolCheckRegionsGauge.Set(dur.Seconds()) - c.setPatrolRegionsDuration(dur) - start = time.Now() } failpoint.Inject("breakPatrol", func() { for !c.IsPatrolRegionChanEmpty() { @@ -219,6 +205,33 @@ func (c *Controller) PatrolRegions() { } } +func (c *Controller) recordPendingPatrolWork() int64 { + pending := c.patrolRegionContext.pendingWorkCount() + c.metrics.patrolRegionPendingWork.Set(float64(pending)) + return pending +} + +func (c *Controller) finishPatrolRound(start time.Time) (time.Time, bool) { + if !c.waitPatrolWorkDone() { + patrolCheckRegionsGauge.Set(0) + c.setPatrolRegionsDuration(0) + return start, false + } + c.patrolRegionScanLimit = calculateScanLimit(c.cluster) + dur := time.Since(start) + patrolCheckRegionsGauge.Set(dur.Seconds()) + c.setPatrolRegionsDuration(dur) + return time.Now(), true +} + +func (c *Controller) waitPatrolWorkDone() bool { + var idle bool + measure(c.metrics.patrolPhaseHistograms[phaseWaitForChannel], func() { + idle = c.patrolRegionContext.waitIdle(c.ctx) + }) + return idle +} + func (c *Controller) updateTickerIfNeeded(ticker *time.Ticker) { // Note: we reset the ticker here to support updating configuration dynamically. newInterval := c.cluster.GetCheckerConfig().GetPatrolRegionInterval() @@ -275,6 +288,14 @@ func newOperatorCandidate(limit operatorLimit, ops ...*operator.Operator) operat return operatorCandidate{ops: ops, limit: limit} } +func newUnlimitedOperatorCandidate(ops ...*operator.Operator) operatorCandidate { + return newOperatorCandidate(operatorLimit{}, ops...) +} + +func (c operatorCandidate) hasOperators() bool { + return len(c.ops) > 0 +} + // GetPatrolRegionsDuration returns the duration of the last patrol region round. func (c *Controller) GetPatrolRegionsDuration() time.Duration { return c.duration.Load().(time.Duration) @@ -320,9 +341,8 @@ func (c *Controller) checkPriorityRegions() { continue } candidate := c.buildOperatorCandidate(region) - ops := candidate.ops // it should skip if region needs to merge - if len(ops) == 0 || ops[0].HasRelatedMergeRegion() { + if !candidate.hasOperators() || candidate.ops[0].HasRelatedMergeRegion() { continue } c.admitWaitingOperators(candidate) @@ -362,7 +382,7 @@ func (c *Controller) mergeLimit() operatorLimit { } } -func (c *Controller) deferByLimit(regionID uint64, limit operatorLimit) { +func (c *Controller) deferRegionByLimit(regionID uint64, limit operatorLimit) { limit.recordLimit() c.pendingProcessedRegions.Put(regionID, nil) } @@ -373,63 +393,27 @@ func (c *Controller) buildOperatorCandidate(region *core.RegionInfo) operatorCan if ops := measureChecker(c.metrics.checkRegionHistograms[jointStateChecker], func() []*operator.Operator { return []*operator.Operator{c.jointStateChecker.Check(region)} }); len(ops) > 0 { - return newOperatorCandidate(operatorLimit{}, ops...) + return newUnlimitedOperatorCandidate(ops...) } if ops := measureChecker(c.metrics.checkRegionHistograms[splitChecker], func() []*operator.Operator { return []*operator.Operator{c.splitChecker.Check(region)} }); len(ops) > 0 { - return newOperatorCandidate(operatorLimit{}, ops...) + return newUnlimitedOperatorCandidate(ops...) } if c.conf.IsPlacementRulesEnabled() { - var candidate operatorCandidate - measure(c.metrics.checkRegionHistograms[ruleChecker], func() { - skipRuleCheck := c.cluster.GetCheckerConfig().IsPlacementRulesCacheEnabled() && - c.cluster.GetRuleManager().IsRegionFitCached(c.cluster, region) - if skipRuleCheck { - // If the fit is fetched from cache, it seems that the region doesn't need check - failpoint.Inject("assertShouldNotCache", func() { - panic("cached shouldn't be used") - }) - ruleCheckerGetCacheCounter.Inc() - } else { - failpoint.Inject("assertShouldCache", func() { - panic("cached should be used") - }) - fit := c.priorityInspector.Inspect(region) - limit := c.replicaLimit(c.ruleChecker.GetType()) - if limit.reached(c.opController) { - c.deferByLimit(region.GetID(), limit) - return - } - if op := c.ruleChecker.CheckWithFit(region, fit); op != nil { - candidate = newOperatorCandidate(limit, op) - } - } - }) - if len(candidate.ops) > 0 { + if candidate := c.checkRuleWithLimit(region); candidate.hasOperators() { return candidate } } else { if ops := measureChecker(c.metrics.checkRegionHistograms[learnerChecker], func() []*operator.Operator { return []*operator.Operator{c.learnerChecker.Check(region)} }); len(ops) > 0 { - return newOperatorCandidate(operatorLimit{}, ops...) + return newUnlimitedOperatorCandidate(ops...) } - var candidate operatorCandidate - measure(c.metrics.checkRegionHistograms[replicaChecker], func() { - if op := c.replicaChecker.Check(region); op != nil { - limit := c.replicaLimit(c.replicaChecker.GetType()) - if limit.reached(c.opController) { - c.deferByLimit(region.GetID(), limit) - return - } - candidate = newOperatorCandidate(limit, op) - } - }) - if len(candidate.ops) > 0 { + if candidate := c.checkReplicaWithLimit(region); candidate.hasOperators() { return candidate } } @@ -441,39 +425,93 @@ func (c *Controller) buildOperatorCandidate(region *core.RegionInfo) operatorCan return operatorCandidate{} } - affinityLimit := c.affinityLimit() - var affinityCandidate operatorCandidate + if candidate := c.checkAffinityWithLimit(region); candidate.hasOperators() { + return candidate + } + + if candidate := c.checkMergeWithLimit(region); candidate.hasOperators() { + return candidate + } + return operatorCandidate{} +} + +func (c *Controller) checkRuleWithLimit(region *core.RegionInfo) operatorCandidate { + var candidate operatorCandidate + measure(c.metrics.checkRegionHistograms[ruleChecker], func() { + skipRuleCheck := c.cluster.GetCheckerConfig().IsPlacementRulesCacheEnabled() && + c.cluster.GetRuleManager().IsRegionFitCached(c.cluster, region) + if skipRuleCheck { + // If the fit is fetched from cache, it seems that the region doesn't need check + failpoint.Inject("assertShouldNotCache", func() { + panic("cached shouldn't be used") + }) + ruleCheckerGetCacheCounter.Inc() + return + } + failpoint.Inject("assertShouldCache", func() { + panic("cached should be used") + }) + fit := c.priorityInspector.Inspect(region) + limit := c.replicaLimit(c.ruleChecker.GetType()) + if limit.reached(c.opController) { + c.deferRegionByLimit(region.GetID(), limit) + return + } + if op := c.ruleChecker.CheckWithFit(region, fit); op != nil { + candidate = newOperatorCandidate(limit, op) + } + }) + return candidate +} + +func (c *Controller) checkReplicaWithLimit(region *core.RegionInfo) operatorCandidate { + var candidate operatorCandidate + measure(c.metrics.checkRegionHistograms[replicaChecker], func() { + op := c.replicaChecker.Check(region) + if op == nil { + return + } + limit := c.replicaLimit(c.replicaChecker.GetType()) + if limit.reached(c.opController) { + c.deferRegionByLimit(region.GetID(), limit) + return + } + candidate = newOperatorCandidate(limit, op) + }) + return candidate +} + +func (c *Controller) checkAffinityWithLimit(region *core.RegionInfo) operatorCandidate { + limit := c.affinityLimit() + var candidate operatorCandidate measure(c.metrics.checkRegionHistograms[affinityChecker], func() { - if affinityLimit.reached(c.opController) { + if limit.reached(c.opController) { if c.affinityChecker.hasAffinityGroups() { - affinityLimit.recordLimit() + limit.recordLimit() } return } if ops := c.affinityChecker.Check(region); len(ops) > 0 { - affinityCandidate = newOperatorCandidate(affinityLimit, ops...) + candidate = newOperatorCandidate(limit, ops...) } }) - if len(affinityCandidate.ops) > 0 { - return affinityCandidate - } + return candidate +} - mergeLimit := c.mergeLimit() - var mergeCandidate operatorCandidate +func (c *Controller) checkMergeWithLimit(region *core.RegionInfo) operatorCandidate { + limit := c.mergeLimit() + var candidate operatorCandidate measure(c.metrics.checkRegionHistograms[mergeChecker], func() { - if mergeLimit.reached(c.opController) { - mergeLimit.recordLimit() + if limit.reached(c.opController) { + limit.recordLimit() return } if ops := c.mergeChecker.Check(region); len(ops) > 0 { // It makes sure that two operators can be added successfully altogether. - mergeCandidate = newOperatorCandidate(mergeLimit, ops...) + candidate = newOperatorCandidate(limit, ops...) } }) - if len(mergeCandidate.ops) > 0 { - return mergeCandidate - } - return operatorCandidate{} + return candidate } func (c *Controller) tryAddOperators(region *core.RegionInfo) { @@ -490,7 +528,7 @@ func (c *Controller) tryAddOperators(region *core.RegionInfo) { return } candidate := c.buildOperatorCandidate(region) - if len(candidate.ops) == 0 { + if !candidate.hasOperators() { return } @@ -505,7 +543,7 @@ func (c *Controller) tryAddOperators(region *core.RegionInfo) { } func (c *Controller) admitWaitingOperators(candidate operatorCandidate) bool { - if len(candidate.ops) == 0 { + if !candidate.hasOperators() { return false } c.operatorMu.Lock() @@ -730,12 +768,12 @@ func (p *PatrolRegionContext) submit(region *core.RegionInfo) { p.regionChan <- region } -func (p *PatrolRegionContext) pendingRegionCount() int64 { +func (p *PatrolRegionContext) pendingWorkCount() int64 { return p.pendingCount.Load() } func (p *PatrolRegionContext) isIdle() bool { - return p.pendingRegionCount() == 0 + return p.pendingWorkCount() == 0 } func (p *PatrolRegionContext) waitIdle(ctx context.Context) bool { diff --git a/pkg/schedule/checker/metrics.go b/pkg/schedule/checker/metrics.go index 0c4c6027171..cb57f477835 100644 --- a/pkg/schedule/checker/metrics.go +++ b/pkg/schedule/checker/metrics.go @@ -58,12 +58,12 @@ var ( Buckets: prometheus.DefBuckets, }, []string{"checker"}) - patrolRegionChannelSize = prometheus.NewGauge( + patrolRegionPendingWork = prometheus.NewGauge( prometheus.GaugeOpts{ Namespace: "pd", Subsystem: "checker", Name: "patrol_region_channel_size", - Help: "Size of patrol region channel.", + Help: "Number of queued and in-flight patrol regions.", }) splitScatterPendingExpiredCounter = prometheus.NewCounterVec( prometheus.CounterOpts{ @@ -100,7 +100,7 @@ func init() { prometheus.MustRegister(patrolCheckRegionsGauge) prometheus.MustRegister(patrolPhaseDuration) prometheus.MustRegister(checkRegionDuration) - prometheus.MustRegister(patrolRegionChannelSize) + prometheus.MustRegister(patrolRegionPendingWork) prometheus.MustRegister(splitScatterPendingExpiredCounter) prometheus.MustRegister(splitScatterPendingDroppedCounter) prometheus.MustRegister(splitScatterPendingGauge) @@ -136,8 +136,8 @@ type checkerControllerMetrics struct { patrolPhaseHistograms map[string]prometheus.Observer // Pre-created histograms for checkers. Key is the checker name. checkRegionHistograms map[string]prometheus.Observer - // Gauge for patrol region channel size. - patrolRegionChannelSize prometheus.Gauge + // Gauge for queued and in-flight patrol regions. + patrolRegionPendingWork prometheus.Gauge } // newCheckerControllerMetrics creates and initializes a new checkerControllerMetrics instance. @@ -165,7 +165,7 @@ func newCheckerControllerMetrics() *checkerControllerMetrics { m := &checkerControllerMetrics{ patrolPhaseHistograms: make(map[string]prometheus.Observer, len(patrolPhases)), checkRegionHistograms: make(map[string]prometheus.Observer, len(checkerTypes)), - patrolRegionChannelSize: patrolRegionChannelSize, // Use the gauge defined at package-level + patrolRegionPendingWork: patrolRegionPendingWork, } for _, phase := range patrolPhases { From 37d9c566f72cd0eac972f9d3b26e6c21d3ba2078 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Thu, 9 Jul 2026 18:24:57 +0800 Subject: [PATCH 8/9] schedule: simplify patrol admission locking Signed-off-by: Ryan Leung --- pkg/schedule/checker/checker_controller.go | 369 +++++------------- .../checker/checker_controller_test.go | 59 --- pkg/schedule/checker/metrics.go | 12 +- 3 files changed, 113 insertions(+), 327 deletions(-) delete mode 100644 pkg/schedule/checker/checker_controller_test.go diff --git a/pkg/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index 4648b7608bb..1dff388447c 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -35,10 +35,8 @@ import ( sche "github.com/tikv/pd/pkg/schedule/core" "github.com/tikv/pd/pkg/schedule/labeler" "github.com/tikv/pd/pkg/schedule/operator" - "github.com/tikv/pd/pkg/schedule/types" "github.com/tikv/pd/pkg/utils/keyutil" "github.com/tikv/pd/pkg/utils/logutil" - "github.com/tikv/pd/pkg/utils/syncutil" ) const ( @@ -86,7 +84,7 @@ type Controller struct { // duration is the duration of the last patrol round. // It's exported, so it should be protected by a mutex. duration atomic.Value // Store as time.Duration - operatorMu syncutil.Mutex + operatorMu sync.Mutex // interval is the config interval of patrol regions. // It's used to update the ticker, so we need to @@ -121,7 +119,6 @@ func NewController(ctx context.Context, cluster sche.CheckerCluster, conf config suspectKeyRanges: cache.NewStringTTL(ctx, time.Minute, 3*time.Minute), patrolRegionContext: &PatrolRegionContext{}, interval: cluster.GetCheckerConfig().GetPatrolRegionInterval(), - workerCount: cluster.GetCheckerConfig().GetPatrolRegionWorkerCount(), patrolRegionScanLimit: calculateScanLimit(cluster), metrics: newCheckerControllerMetrics(), } @@ -151,14 +148,22 @@ func (c *Controller) PatrolRegions() { failpoint.Continue() }) c.updateTickerIfNeeded(ticker) - if c.recordPendingPatrolWork() > 0 { - continue - } c.updatePatrolWorkersIfNeeded() if c.cluster.IsSchedulingHalted() { + for len(c.patrolRegionContext.regionChan) > 0 { + <-c.patrolRegionContext.regionChan + } log.Debug("skip patrol regions due to scheduling is halted") continue } + c.metrics.patrolRegionChannelSize.Set(float64(len(c.patrolRegionContext.regionChan))) + + // wait for the regionChan to be drained + waitDrainChanel := time.Now() + if len(c.patrolRegionContext.regionChan) > 0 { + c.metrics.patrolPhaseHistograms[phaseWaitForChannel].Observe(time.Since(waitDrainChanel).Seconds()) + continue + } // Check priority regions and pending processed regions first. measure(c.metrics.patrolPhaseHistograms[phaseCheckPriority], func() { @@ -185,11 +190,13 @@ func (c *Controller) PatrolRegions() { }) // When the key is nil, it means that the scan is finished. if len(key) == 0 { - var completed bool - start, completed = c.finishPatrolRound(start) - if !completed { - return - } + // update the scan limit. + c.patrolRegionScanLimit = calculateScanLimit(c.cluster) + // update the metrics. + dur := time.Since(start) + patrolCheckRegionsGauge.Set(dur.Seconds()) + c.setPatrolRegionsDuration(dur) + start = time.Now() } failpoint.Inject("breakPatrol", func() { for !c.IsPatrolRegionChanEmpty() { @@ -205,33 +212,6 @@ func (c *Controller) PatrolRegions() { } } -func (c *Controller) recordPendingPatrolWork() int64 { - pending := c.patrolRegionContext.pendingWorkCount() - c.metrics.patrolRegionPendingWork.Set(float64(pending)) - return pending -} - -func (c *Controller) finishPatrolRound(start time.Time) (time.Time, bool) { - if !c.waitPatrolWorkDone() { - patrolCheckRegionsGauge.Set(0) - c.setPatrolRegionsDuration(0) - return start, false - } - c.patrolRegionScanLimit = calculateScanLimit(c.cluster) - dur := time.Since(start) - patrolCheckRegionsGauge.Set(dur.Seconds()) - c.setPatrolRegionsDuration(dur) - return time.Now(), true -} - -func (c *Controller) waitPatrolWorkDone() bool { - var idle bool - measure(c.metrics.patrolPhaseHistograms[phaseWaitForChannel], func() { - idle = c.patrolRegionContext.waitIdle(c.ctx) - }) - return idle -} - func (c *Controller) updateTickerIfNeeded(ticker *time.Ticker) { // Note: we reset the ticker here to support updating configuration dynamically. newInterval := c.cluster.GetCheckerConfig().GetPatrolRegionInterval() @@ -245,9 +225,6 @@ func (c *Controller) updateTickerIfNeeded(ticker *time.Ticker) { func (c *Controller) updatePatrolWorkersIfNeeded() { newWorkersCount := c.cluster.GetCheckerConfig().GetPatrolRegionWorkerCount() if c.workerCount != newWorkersCount { - if !c.patrolRegionContext.isIdle() { - return - } oldWorkersCount := c.workerCount c.workerCount = newWorkersCount // Stop the old workers and start the new workers. @@ -261,41 +238,6 @@ func (c *Controller) updatePatrolWorkersIfNeeded() { } } -// operatorLimit is rechecked while admitting a candidate into the operator queue. -type operatorLimit struct { - kind operator.OpKind - checkerType types.CheckerSchedulerType - getLimit func() uint64 -} - -func (l operatorLimit) reached(opController *operator.Controller) bool { - return l.getLimit != nil && opController.OperatorCount(l.kind) >= l.getLimit() -} - -func (l operatorLimit) recordLimit() { - if l.getLimit != nil { - operator.IncOperatorLimitCounter(l.checkerType, l.kind) - } -} - -// operatorCandidate is a checker result that still needs operator-controller admission. -type operatorCandidate struct { - ops []*operator.Operator - limit operatorLimit -} - -func newOperatorCandidate(limit operatorLimit, ops ...*operator.Operator) operatorCandidate { - return operatorCandidate{ops: ops, limit: limit} -} - -func newUnlimitedOperatorCandidate(ops ...*operator.Operator) operatorCandidate { - return newOperatorCandidate(operatorLimit{}, ops...) -} - -func (c operatorCandidate) hasOperators() bool { - return len(c.ops) > 0 -} - // GetPatrolRegionsDuration returns the duration of the last patrol region round. func (c *Controller) GetPatrolRegionsDuration() time.Duration { return c.duration.Load().(time.Duration) @@ -314,7 +256,7 @@ func (c *Controller) checkRegions(startKey []byte) (key []byte, regions []*core. } for _, region := range regions { - c.patrolRegionContext.submit(region) + c.patrolRegionContext.regionChan <- region key = region.GetEndKey() } return @@ -340,81 +282,91 @@ func (c *Controller) checkPriorityRegions() { removes = append(removes, id) continue } - candidate := c.buildOperatorCandidate(region) - // it should skip if region needs to merge - if !candidate.hasOperators() || candidate.ops[0].HasRelatedMergeRegion() { - continue - } - c.admitWaitingOperators(candidate) + c.addPriorityOperator(region) } for _, v := range removes { c.RemovePriorityRegions(v) } } -// CheckRegion will check the region and add a new operator if needed. -// The function is exposed for test purpose. -func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { - return c.buildOperatorCandidate(region).ops -} - -func (c *Controller) replicaLimit(checkerType types.CheckerSchedulerType) operatorLimit { - return operatorLimit{ - kind: operator.OpReplica, - checkerType: checkerType, - getLimit: c.conf.GetReplicaScheduleLimit, +func (c *Controller) addPriorityOperator(region *core.RegionInfo) { + c.operatorMu.Lock() + defer c.operatorMu.Unlock() + if c.cluster.IsSchedulingHalted() { + return } -} - -func (c *Controller) affinityLimit() operatorLimit { - return operatorLimit{ - kind: operator.OpAffinity, - checkerType: c.affinityChecker.GetType(), - getLimit: c.conf.GetAffinityScheduleLimit, + ops := c.CheckRegion(region) + // It should skip if region needs to merge. + if len(ops) == 0 || ops[0].HasRelatedMergeRegion() { + return } -} - -func (c *Controller) mergeLimit() operatorLimit { - return operatorLimit{ - kind: operator.OpMerge, - checkerType: c.mergeChecker.GetType(), - getLimit: c.conf.GetMergeScheduleLimit, + if !c.opController.ExceedStoreLimit(ops...) { + c.opController.AddWaitingOperator(ops...) } } -func (c *Controller) deferRegionByLimit(regionID uint64, limit operatorLimit) { - limit.recordLimit() - c.pendingProcessedRegions.Put(regionID, nil) -} - -func (c *Controller) buildOperatorCandidate(region *core.RegionInfo) operatorCandidate { +// CheckRegion will check the region and add a new operator if needed. +// The function is exposed for test purpose. +func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { // If PD has restarted, it needs to check learners added before and promote them. // Don't check isRaftLearnerEnabled cause it maybe disable learner feature but there are still some learners to promote. + opController := c.opController + if ops := measureChecker(c.metrics.checkRegionHistograms[jointStateChecker], func() []*operator.Operator { return []*operator.Operator{c.jointStateChecker.Check(region)} }); len(ops) > 0 { - return newUnlimitedOperatorCandidate(ops...) + return ops } if ops := measureChecker(c.metrics.checkRegionHistograms[splitChecker], func() []*operator.Operator { return []*operator.Operator{c.splitChecker.Check(region)} }); len(ops) > 0 { - return newUnlimitedOperatorCandidate(ops...) + return ops } if c.conf.IsPlacementRulesEnabled() { - if candidate := c.checkRuleWithLimit(region); candidate.hasOperators() { - return candidate + if ops := measureChecker(c.metrics.checkRegionHistograms[ruleChecker], func() []*operator.Operator { + skipRuleCheck := c.cluster.GetCheckerConfig().IsPlacementRulesCacheEnabled() && + c.cluster.GetRuleManager().IsRegionFitCached(c.cluster, region) + if skipRuleCheck { + // If the fit is fetched from cache, it seems that the region doesn't need check + failpoint.Inject("assertShouldNotCache", func() { + panic("cached shouldn't be used") + }) + ruleCheckerGetCacheCounter.Inc() + } else { + failpoint.Inject("assertShouldCache", func() { + panic("cached should be used") + }) + fit := c.priorityInspector.Inspect(region) + if opController.OperatorCount(operator.OpReplica) < c.conf.GetReplicaScheduleLimit() { + return []*operator.Operator{c.ruleChecker.CheckWithFit(region, fit)} + } + operator.IncOperatorLimitCounter(c.ruleChecker.GetType(), operator.OpReplica) + c.pendingProcessedRegions.Put(region.GetID(), nil) + } + return nil + }); len(ops) > 0 { + return ops } } else { if ops := measureChecker(c.metrics.checkRegionHistograms[learnerChecker], func() []*operator.Operator { return []*operator.Operator{c.learnerChecker.Check(region)} }); len(ops) > 0 { - return newUnlimitedOperatorCandidate(ops...) + return ops } - if candidate := c.checkReplicaWithLimit(region); candidate.hasOperators() { - return candidate + if ops := measureChecker(c.metrics.checkRegionHistograms[replicaChecker], func() []*operator.Operator { + if op := c.replicaChecker.Check(region); op != nil { + if opController.OperatorCount(operator.OpReplica) < c.conf.GetReplicaScheduleLimit() { + return []*operator.Operator{op} + } + operator.IncOperatorLimitCounter(c.replicaChecker.GetType(), operator.OpReplica) + c.pendingProcessedRegions.Put(region.GetID(), nil) + } + return nil + }); len(ops) > 0 { + return ops } } // skip the joint checker, split checker and rule checker when region label is set to "schedule=deny". @@ -422,96 +374,33 @@ func (c *Controller) buildOperatorCandidate(region *core.RegionInfo) operatorCan l := c.cluster.GetRegionLabeler() if l.ScheduleDisabled(region) { denyCheckersByLabelerCounter.Inc() - return operatorCandidate{} - } - - if candidate := c.checkAffinityWithLimit(region); candidate.hasOperators() { - return candidate - } - - if candidate := c.checkMergeWithLimit(region); candidate.hasOperators() { - return candidate + return nil } - return operatorCandidate{} -} - -func (c *Controller) checkRuleWithLimit(region *core.RegionInfo) operatorCandidate { - var candidate operatorCandidate - measure(c.metrics.checkRegionHistograms[ruleChecker], func() { - skipRuleCheck := c.cluster.GetCheckerConfig().IsPlacementRulesCacheEnabled() && - c.cluster.GetRuleManager().IsRegionFitCached(c.cluster, region) - if skipRuleCheck { - // If the fit is fetched from cache, it seems that the region doesn't need check - failpoint.Inject("assertShouldNotCache", func() { - panic("cached shouldn't be used") - }) - ruleCheckerGetCacheCounter.Inc() - return - } - failpoint.Inject("assertShouldCache", func() { - panic("cached should be used") - }) - fit := c.priorityInspector.Inspect(region) - limit := c.replicaLimit(c.ruleChecker.GetType()) - if limit.reached(c.opController) { - c.deferRegionByLimit(region.GetID(), limit) - return - } - if op := c.ruleChecker.CheckWithFit(region, fit); op != nil { - candidate = newOperatorCandidate(limit, op) - } - }) - return candidate -} - -func (c *Controller) checkReplicaWithLimit(region *core.RegionInfo) operatorCandidate { - var candidate operatorCandidate - measure(c.metrics.checkRegionHistograms[replicaChecker], func() { - op := c.replicaChecker.Check(region) - if op == nil { - return - } - limit := c.replicaLimit(c.replicaChecker.GetType()) - if limit.reached(c.opController) { - c.deferRegionByLimit(region.GetID(), limit) - return - } - candidate = newOperatorCandidate(limit, op) - }) - return candidate -} -func (c *Controller) checkAffinityWithLimit(region *core.RegionInfo) operatorCandidate { - limit := c.affinityLimit() - var candidate operatorCandidate - measure(c.metrics.checkRegionHistograms[affinityChecker], func() { - if limit.reached(c.opController) { - if c.affinityChecker.hasAffinityGroups() { - limit.recordLimit() - } - return + if ops := measureChecker(c.metrics.checkRegionHistograms[affinityChecker], func() []*operator.Operator { + if opController.OperatorCount(operator.OpAffinity) < c.conf.GetAffinityScheduleLimit() { + // It makes sure that two affinity merge operators can be added successfully altogether. + return c.affinityChecker.Check(region) } - if ops := c.affinityChecker.Check(region); len(ops) > 0 { - candidate = newOperatorCandidate(limit, ops...) + if c.affinityChecker.hasAffinityGroups() { + operator.IncOperatorLimitCounter(c.affinityChecker.GetType(), operator.OpAffinity) } - }) - return candidate -} + return nil + }); len(ops) > 0 { + return ops + } -func (c *Controller) checkMergeWithLimit(region *core.RegionInfo) operatorCandidate { - limit := c.mergeLimit() - var candidate operatorCandidate - measure(c.metrics.checkRegionHistograms[mergeChecker], func() { - if limit.reached(c.opController) { - limit.recordLimit() - return - } - if ops := c.mergeChecker.Check(region); len(ops) > 0 { + if ops := measureChecker(c.metrics.checkRegionHistograms[mergeChecker], func() []*operator.Operator { + if opController.OperatorCount(operator.OpMerge) < c.conf.GetMergeScheduleLimit() { // It makes sure that two operators can be added successfully altogether. - candidate = newOperatorCandidate(limit, ops...) + return c.mergeChecker.Check(region) } - }) - return candidate + operator.IncOperatorLimitCounter(c.mergeChecker.GetType(), operator.OpMerge) + return nil + }); len(ops) > 0 { + return ops + } + return nil } func (c *Controller) tryAddOperators(region *core.RegionInfo) { @@ -519,48 +408,33 @@ func (c *Controller) tryAddOperators(region *core.RegionInfo) { // the region could be recent split, continue to wait. return } + id := region.GetID() + if c.opController.GetOperator(id) != nil { + c.RemovePendingProcessedRegion(id) + return + } + c.operatorMu.Lock() + defer c.operatorMu.Unlock() if c.cluster.IsSchedulingHalted() { return } - id := region.GetID() if c.opController.GetOperator(id) != nil { c.RemovePendingProcessedRegion(id) return } - candidate := c.buildOperatorCandidate(region) - if !candidate.hasOperators() { + ops := c.CheckRegion(region) + if len(ops) == 0 { return } - if c.admitWaitingOperators(candidate) { + if !c.opController.ExceedStoreLimit(ops...) { + c.opController.AddWaitingOperator(ops...) c.RemovePendingProcessedRegion(id) } else { - if c.cluster.IsSchedulingHalted() { - return - } c.AddPendingProcessedRegions(true, id) } } -func (c *Controller) admitWaitingOperators(candidate operatorCandidate) bool { - if !candidate.hasOperators() { - return false - } - c.operatorMu.Lock() - defer c.operatorMu.Unlock() - if c.cluster.IsSchedulingHalted() { - return false - } - if candidate.limit.reached(c.opController) { - candidate.limit.recordLimit() - return false - } - if c.opController.ExceedStoreLimit(candidate.ops...) { - return false - } - return c.opController.AddWaitingOperator(candidate.ops...) > 0 -} - // GetMergeChecker returns the merge checker. func (c *Controller) GetMergeChecker() *MergeChecker { return c.mergeChecker @@ -711,7 +585,7 @@ func (c *Controller) IsPatrolRegionChanEmpty() bool { if c.patrolRegionContext == nil { return true } - return c.patrolRegionContext.isIdle() + return len(c.patrolRegionContext.regionChan) == 0 } // PatrolRegionContext is used to store the context of patrol regions. @@ -720,13 +594,11 @@ type PatrolRegionContext struct { workersCancel context.CancelFunc regionChan chan *core.RegionInfo wg sync.WaitGroup - pendingCount atomic.Int64 } func (p *PatrolRegionContext) init(ctx context.Context) { p.regionChan = make(chan *core.RegionInfo, patrolRegionChanLen) p.workersCtx, p.workersCancel = context.WithCancel(ctx) - p.pendingCount.Store(0) } func (p *PatrolRegionContext) stop() { @@ -750,10 +622,7 @@ func (p *PatrolRegionContext) startPatrolRegionWorkers(c *Controller) { log.Debug("region channel is closed", zap.Int("worker-id", i)) return } - func() { - defer p.pendingCount.Add(-1) - c.tryAddOperators(region) - }() + c.tryAddOperators(region) case <-p.workersCtx.Done(): log.Debug("region worker is closed", zap.Int("worker-id", i)) return @@ -763,30 +632,6 @@ func (p *PatrolRegionContext) startPatrolRegionWorkers(c *Controller) { } } -func (p *PatrolRegionContext) submit(region *core.RegionInfo) { - p.pendingCount.Add(1) - p.regionChan <- region -} - -func (p *PatrolRegionContext) pendingWorkCount() int64 { - return p.pendingCount.Load() -} - -func (p *PatrolRegionContext) isIdle() bool { - return p.pendingWorkCount() == 0 -} - -func (p *PatrolRegionContext) waitIdle(ctx context.Context) bool { - for !p.isIdle() { - select { - case <-ctx.Done(): - return false - case <-time.After(time.Millisecond): - } - } - return true -} - // GetPatrolRegionScanLimit returns the limit of regions to scan. // It only used for test. func (c *Controller) GetPatrolRegionScanLimit() int { diff --git a/pkg/schedule/checker/checker_controller_test.go b/pkg/schedule/checker/checker_controller_test.go deleted file mode 100644 index 2851bcd60fe..00000000000 --- a/pkg/schedule/checker/checker_controller_test.go +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright 2026 TiKV Project Authors. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package checker - -import ( - "context" - "testing" - - "github.com/stretchr/testify/require" - - "github.com/tikv/pd/pkg/mock/mockcluster" - "github.com/tikv/pd/pkg/mock/mockconfig" - "github.com/tikv/pd/pkg/schedule/hbstream" - "github.com/tikv/pd/pkg/schedule/operator" - "github.com/tikv/pd/pkg/schedule/types" -) - -func TestAddWaitingOperatorsSkipsSchedulingHalted(t *testing.T) { - re := require.New(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - opt := mockconfig.NewTestOptions() - tc := mockcluster.NewCluster(ctx, opt) - for storeID := uint64(1); storeID <= 3; storeID++ { - tc.AddRegionStore(storeID, 0) - } - region := tc.AddLeaderRegion(1, 1, 2) - - stream := hbstream.NewTestHeartbeatStreams(ctx, tc, false) - defer stream.Close() - oc := operator.NewController(ctx, tc.GetBasicCluster(), tc.GetSharedConfig(), stream) - controller := NewController(ctx, tc, tc.GetCheckerConfig(), oc) - candidate := newOperatorCandidate( - controller.replicaLimit(types.ReplicaChecker), - operator.NewTestOperator(region.GetID(), region.GetRegionEpoch(), operator.OpReplica), - ) - - tc.SetHaltScheduling(true, "test") - re.False(controller.admitWaitingOperators(candidate)) - re.Empty(oc.GetOperators()) - re.Empty(oc.GetWaitingOperators()) - - tc.SetHaltScheduling(false, "test") - re.True(controller.admitWaitingOperators(candidate)) - re.Len(oc.GetOperators(), 1) -} diff --git a/pkg/schedule/checker/metrics.go b/pkg/schedule/checker/metrics.go index cb57f477835..0c4c6027171 100644 --- a/pkg/schedule/checker/metrics.go +++ b/pkg/schedule/checker/metrics.go @@ -58,12 +58,12 @@ var ( Buckets: prometheus.DefBuckets, }, []string{"checker"}) - patrolRegionPendingWork = prometheus.NewGauge( + patrolRegionChannelSize = prometheus.NewGauge( prometheus.GaugeOpts{ Namespace: "pd", Subsystem: "checker", Name: "patrol_region_channel_size", - Help: "Number of queued and in-flight patrol regions.", + Help: "Size of patrol region channel.", }) splitScatterPendingExpiredCounter = prometheus.NewCounterVec( prometheus.CounterOpts{ @@ -100,7 +100,7 @@ func init() { prometheus.MustRegister(patrolCheckRegionsGauge) prometheus.MustRegister(patrolPhaseDuration) prometheus.MustRegister(checkRegionDuration) - prometheus.MustRegister(patrolRegionPendingWork) + prometheus.MustRegister(patrolRegionChannelSize) prometheus.MustRegister(splitScatterPendingExpiredCounter) prometheus.MustRegister(splitScatterPendingDroppedCounter) prometheus.MustRegister(splitScatterPendingGauge) @@ -136,8 +136,8 @@ type checkerControllerMetrics struct { patrolPhaseHistograms map[string]prometheus.Observer // Pre-created histograms for checkers. Key is the checker name. checkRegionHistograms map[string]prometheus.Observer - // Gauge for queued and in-flight patrol regions. - patrolRegionPendingWork prometheus.Gauge + // Gauge for patrol region channel size. + patrolRegionChannelSize prometheus.Gauge } // newCheckerControllerMetrics creates and initializes a new checkerControllerMetrics instance. @@ -165,7 +165,7 @@ func newCheckerControllerMetrics() *checkerControllerMetrics { m := &checkerControllerMetrics{ patrolPhaseHistograms: make(map[string]prometheus.Observer, len(patrolPhases)), checkRegionHistograms: make(map[string]prometheus.Observer, len(checkerTypes)), - patrolRegionPendingWork: patrolRegionPendingWork, + patrolRegionChannelSize: patrolRegionChannelSize, // Use the gauge defined at package-level } for _, phase := range patrolPhases { From 8e7ec80854c368f3a0e5857417ca949ccacdf890 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Fri, 10 Jul 2026 12:00:30 +0800 Subject: [PATCH 9/9] schedule: account for merge candidate count Signed-off-by: Ryan Leung --- pkg/schedule/checker/checker_controller.go | 24 +++++++--- pkg/schedule/checker/merge_checker_test.go | 56 ++++++++++++++++++++++ 2 files changed, 74 insertions(+), 6 deletions(-) diff --git a/pkg/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index 1dff388447c..dd8453c805d 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -378,9 +378,13 @@ func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { } if ops := measureChecker(c.metrics.checkRegionHistograms[affinityChecker], func() []*operator.Operator { - if opController.OperatorCount(operator.OpAffinity) < c.conf.GetAffinityScheduleLimit() { - // It makes sure that two affinity merge operators can be added successfully altogether. - return c.affinityChecker.Check(region) + current := opController.OperatorCount(operator.OpAffinity) + limit := c.conf.GetAffinityScheduleLimit() + if current < limit { + ops := c.affinityChecker.Check(region) + if hasOperatorCapacity(current, limit, uint64(len(ops))) { + return ops + } } if c.affinityChecker.hasAffinityGroups() { operator.IncOperatorLimitCounter(c.affinityChecker.GetType(), operator.OpAffinity) @@ -391,9 +395,13 @@ func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { } if ops := measureChecker(c.metrics.checkRegionHistograms[mergeChecker], func() []*operator.Operator { - if opController.OperatorCount(operator.OpMerge) < c.conf.GetMergeScheduleLimit() { - // It makes sure that two operators can be added successfully altogether. - return c.mergeChecker.Check(region) + current := opController.OperatorCount(operator.OpMerge) + limit := c.conf.GetMergeScheduleLimit() + if current < limit { + ops := c.mergeChecker.Check(region) + if hasOperatorCapacity(current, limit, uint64(len(ops))) { + return ops + } } operator.IncOperatorLimitCounter(c.mergeChecker.GetType(), operator.OpMerge) return nil @@ -403,6 +411,10 @@ func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { return nil } +func hasOperatorCapacity(current, limit, candidateCount uint64) bool { + return current < limit && candidateCount <= limit-current +} + func (c *Controller) tryAddOperators(region *core.RegionInfo) { if region == nil { // the region could be recent split, continue to wait. diff --git a/pkg/schedule/checker/merge_checker_test.go b/pkg/schedule/checker/merge_checker_test.go index 80a72238ee4..a1ff1386c95 100644 --- a/pkg/schedule/checker/merge_checker_test.go +++ b/pkg/schedule/checker/merge_checker_test.go @@ -20,6 +20,7 @@ import ( "testing" "time" + "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" "go.uber.org/goleak" @@ -56,6 +57,61 @@ func TestMergeCheckerTestSuite(t *testing.T) { suite.Run(t, new(mergeCheckerTestSuite)) } +func TestTryAddOperatorsRejectsMergePairPastLimit(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + opt := mockconfig.NewTestOptions() + tc := mockcluster.NewCluster(ctx, opt) + tc.SetMergeScheduleLimit(3) + tc.SetSplitMergeInterval(0) + for storeID := uint64(1); storeID <= 3; storeID++ { + tc.AddRegionStore(storeID, 0) + } + tc.AddLeaderRegionWithRange(1, "", "a", 1, 2, 3) + tc.AddLeaderRegionWithRange(2, "a", "b", 1, 2, 3) + tc.AddLeaderRegionWithRange(3, "b", "c", 1, 2, 3) + tc.AddLeaderRegionWithRange(4, "c", "d", 1, 2, 3) + + stream := hbstream.NewTestHeartbeatStreams(ctx, tc, false) + defer stream.Close() + oc := operator.NewController(ctx, tc.GetBasicCluster(), tc.GetSharedConfig(), stream) + controller := NewController(ctx, tc, tc.GetCheckerConfig(), oc) + + controller.tryAddOperators(tc.GetRegion(1)) + re.Equal(uint64(2), oc.OperatorCount(operator.OpMerge)) + + controller.tryAddOperators(tc.GetRegion(3)) + re.LessOrEqual(oc.OperatorCount(operator.OpMerge), uint64(3)) +} + +func TestHasOperatorCapacity(t *testing.T) { + testCases := []struct { + name string + current uint64 + limit uint64 + candidateCount uint64 + expected bool + }{ + {name: "no capacity", current: 0, limit: 0, candidateCount: 1, expected: false}, + {name: "single candidate fits", current: 2, limit: 3, candidateCount: 1, expected: true}, + {name: "pair fits", current: 2, limit: 4, candidateCount: 2, expected: true}, + {name: "pair exceeds remaining capacity", current: 2, limit: 3, candidateCount: 2, expected: false}, + {name: "current exceeds limit", current: 4, limit: 3, candidateCount: 1, expected: false}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + require.Equal(t, testCase.expected, hasOperatorCapacity( + testCase.current, + testCase.limit, + testCase.candidateCount, + )) + }) + } +} + func (suite *mergeCheckerTestSuite) SetupTest() { cfg := mockconfig.NewTestOptions() gcInterval = 100 * time.Millisecond