diff --git a/pkg/mcs/scheduling/server/apis/v1/api.go b/pkg/mcs/scheduling/server/apis/v1/api.go index 52bcc399509..68685508582 100644 --- a/pkg/mcs/scheduling/server/apis/v1/api.go +++ b/pkg/mcs/scheduling/server/apis/v1/api.go @@ -383,7 +383,7 @@ func deleteAllRegionCache(c *gin.Context) { c.String(http.StatusInternalServerError, errs.ErrNotBootstrapped.GenWithStackByArgs().Error()) return } - cluster.ResetRegionCache() + cluster.ResetPreparedAndResetRegionCache() c.String(http.StatusOK, "All regions are removed from server cache.") } diff --git a/pkg/mcs/scheduling/server/cluster.go b/pkg/mcs/scheduling/server/cluster.go index 7b846a241ae..139a48457be 100644 --- a/pkg/mcs/scheduling/server/cluster.go +++ b/pkg/mcs/scheduling/server/cluster.go @@ -939,6 +939,17 @@ func (c *Cluster) SetPrepared() { c.coordinator.GetPrepareChecker().SetPrepared() } +// ResetPrepared reset the prepare checker. +func (c *Cluster) ResetPrepared() { + c.coordinator.GetPrepareChecker().ResetPrepared() +} + +// ResetPreparedAndResetRegionCache atomically pauses scheduling and clears the +// region cache. +func (c *Cluster) ResetPreparedAndResetRegionCache() { + c.coordinator.GetPrepareChecker().ResetPreparedAndRun(c.ResetRegionCache) +} + // IsSchedulingHalted returns whether the scheduling is halted. // Currently, the microservice scheduling is halted when: // - The `HaltScheduling` persist option is set to true. diff --git a/pkg/mcs/scheduling/server/server.go b/pkg/mcs/scheduling/server/server.go index ce621211a46..bd696c1e85b 100644 --- a/pkg/mcs/scheduling/server/server.go +++ b/pkg/mcs/scheduling/server/server.go @@ -222,7 +222,7 @@ func (s *Server) updatePDMemberLoop() { if cluster != nil { if cluster.SwitchPDLeader(pdpb.NewPDClient(cc)) { if status.Leader != curLeader { - log.Info("switch PD leader", zap.String("leader-id", strconv.FormatUint(ep.ID, 16)), zap.String("endpoint", ep.ClientURLs[0])) + log.Info("switch PD leader", zap.String("current-leader", strconv.FormatUint(curLeader, 16)), zap.String("new-leader-id", strconv.FormatUint(ep.ID, 16)), zap.String("endpoint", ep.ClientURLs[0])) } curLeader = ep.ID break diff --git a/pkg/mock/mockcluster/mockcluster.go b/pkg/mock/mockcluster/mockcluster.go index 4705876c6a8..9b5fd603e3a 100644 --- a/pkg/mock/mockcluster/mockcluster.go +++ b/pkg/mock/mockcluster/mockcluster.go @@ -91,6 +91,11 @@ func NewCluster(ctx context.Context, opts *config.PersistOptions) *Cluster { return c } +// Context returns the cluster context. +func (mc *Cluster) Context() context.Context { + return mc.ctx +} + // GetStoreConfig returns the store config. func (mc *Cluster) GetStoreConfig() sc.StoreConfigProvider { return mc.PersistOptions.GetStoreConfig() @@ -937,3 +942,12 @@ func (mc *Cluster) ObserveRegionsStats() { storeIDs, writeBytesRates, writeKeysRates := mc.GetStoresWriteRate() mc.HotStat.ObserveRegionsStats(storeIDs, writeBytesRates, writeKeysRates) } + +// ResetPrepared mocks method. +func (*Cluster) ResetPrepared() {} + +// ResetPreparedAndResetRegionCache mocks method. +func (mc *Cluster) ResetPreparedAndResetRegionCache(context.Context) error { + mc.ResetRegionCache() + return nil +} diff --git a/pkg/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index 7da5ae7bb49..4001a50ba54 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -35,6 +35,7 @@ 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/preparecheck" "github.com/tikv/pd/pkg/utils/keyutil" "github.com/tikv/pd/pkg/utils/logutil" ) @@ -94,13 +95,20 @@ type Controller struct { // patrolRegionScanLimit is the limit of regions to scan. // It is calculated by the number of regions. patrolRegionScanLimit int + prepareChecker *preparecheck.Checker metrics *checkerControllerMetrics } // NewController create a new Controller. -func NewController(ctx context.Context, cluster sche.CheckerCluster, conf config.CheckerConfigProvider, opController *operator.Controller) *Controller { +func NewController(ctx context.Context, cluster sche.CheckerCluster, conf config.CheckerConfigProvider, opController *operator.Controller, prepareCheckers ...*preparecheck.Checker) *Controller { pendingProcessedRegions := cache.NewIDTTL(ctx, time.Minute, 3*time.Minute) ruleManager := cluster.GetRuleManager() + prepareChecker := preparecheck.NewChecker(cluster.GetPrepareRegionCount) + if len(prepareCheckers) > 0 && prepareCheckers[0] != nil { + prepareChecker = prepareCheckers[0] + } else { + prepareChecker.SetPrepared() + } c := &Controller{ ctx: ctx, cluster: cluster, @@ -119,6 +127,7 @@ func NewController(ctx context.Context, cluster sche.CheckerCluster, conf config patrolRegionContext: &PatrolRegionContext{}, interval: cluster.GetCheckerConfig().GetPatrolRegionInterval(), patrolRegionScanLimit: calculateScanLimit(cluster), + prepareChecker: prepareChecker, metrics: newCheckerControllerMetrics(), } c.splitScatter = newSplitScatterController(ctx, cluster, opController, c.AddPendingProcessedRegions) @@ -148,6 +157,9 @@ func (c *Controller) PatrolRegions() { }) c.updateTickerIfNeeded(ticker) c.updatePatrolWorkersIfNeeded() + if !c.prepareChecker.IsPrepared() { + continue + } if c.cluster.IsSchedulingHalted() { for len(c.patrolRegionContext.regionChan) > 0 { <-c.patrolRegionContext.regionChan @@ -174,7 +186,7 @@ func (c *Controller) PatrolRegions() { }) measure(c.metrics.patrolPhaseHistograms[phaseDispatchSplitScatter], func() { - c.splitScatter.dispatchSplitScatterRegions() + c.dispatchSplitScatterRegions() }) measure(c.metrics.patrolPhaseHistograms[phaseScanRegions], func() { @@ -211,6 +223,10 @@ func (c *Controller) PatrolRegions() { } } +func (c *Controller) dispatchSplitScatterRegions() { + c.prepareChecker.RunIfPrepared(c.splitScatter.dispatchSplitScatterRegions) +} + func (c *Controller) updateTickerIfNeeded(ticker *time.Ticker) { // Note: we reset the ticker here to support updating configuration dynamically. newInterval := c.cluster.GetCheckerConfig().GetPatrolRegionInterval() @@ -272,27 +288,29 @@ func (c *Controller) checkPendingProcessedRegions() { // checkPriorityRegions checks priority regions func (c *Controller) checkPriorityRegions() { - items := c.GetPriorityRegions() - removes := make([]uint64, 0) - priorityListGauge.Set(float64(len(items))) - for _, id := range items { - region := c.cluster.GetRegion(id) - if region == nil { - removes = append(removes, id) - continue - } - ops := c.CheckRegion(region) - // it should skip if region needs to merge - if len(ops) == 0 || ops[0].HasRelatedMergeRegion() { - continue + c.prepareChecker.RunIfPrepared(func() { + items := c.GetPriorityRegions() + removes := make([]uint64, 0) + priorityListGauge.Set(float64(len(items))) + for _, id := range items { + region := c.cluster.GetRegion(id) + if region == nil { + removes = append(removes, id) + continue + } + ops := c.CheckRegion(region) + // 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...) + } } - if !c.opController.ExceedStoreLimit(ops...) { - c.opController.AddWaitingOperator(ops...) + for _, v := range removes { + c.RemovePriorityRegions(v) } - } - for _, v := range removes { - c.RemovePriorityRegions(v) - } + }) } // CheckRegion will check the region and add a new operator if needed. @@ -394,6 +412,12 @@ func (c *Controller) CheckRegion(region *core.RegionInfo) []*operator.Operator { } func (c *Controller) tryAddOperators(region *core.RegionInfo) { + c.prepareChecker.RunIfPrepared(func() { + c.tryAddOperatorsLocked(region) + }) +} + +func (c *Controller) tryAddOperatorsLocked(region *core.RegionInfo) { if region == nil { // the region could be recent split, continue to wait. return diff --git a/pkg/schedule/checker/checker_controller_test.go b/pkg/schedule/checker/checker_controller_test.go new file mode 100644 index 00000000000..891ca069eff --- /dev/null +++ b/pkg/schedule/checker/checker_controller_test.go @@ -0,0 +1,57 @@ +// 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/preparecheck" +) + +func TestTryAddOperatorsRespectsPrepareChecker(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + opts := mockconfig.NewTestOptions() + cluster := mockcluster.NewCluster(ctx, opts) + cluster.SetEnablePlacementRules(false) + cluster.SetMaxReplicas(3) + cluster.AddRegionStore(1, 1) + cluster.AddRegionStore(2, 0) + cluster.AddRegionStore(3, 0) + region := cluster.AddLeaderRegion(1, 1) + + stream := hbstream.NewTestHeartbeatStreams(ctx, cluster, false /* no need to run */) + defer stream.Close() + opController := operator.NewController(ctx, cluster.GetBasicCluster(), cluster.GetSharedConfig(), stream) + prepareChecker := preparecheck.NewChecker(cluster.GetPrepareRegionCount) + controller := NewController(ctx, cluster, cluster.GetCheckerConfig(), opController, prepareChecker) + + controller.tryAddOperators(region) + re.Empty(opController.GetWaitingOperators()) + re.Empty(opController.GetOperators()) + + prepareChecker.SetPrepared() + controller.tryAddOperators(region) + re.NotEmpty(append(opController.GetWaitingOperators(), opController.GetOperators()...)) +} diff --git a/pkg/schedule/checker/split_scatter_test.go b/pkg/schedule/checker/split_scatter_test.go index 0d6076716f1..b2d4eecb9a2 100644 --- a/pkg/schedule/checker/split_scatter_test.go +++ b/pkg/schedule/checker/split_scatter_test.go @@ -54,10 +54,6 @@ func (c *Controller) collectTopPendingSplitScatter(limit int) []splitScatterPend return c.splitScatter.collectTopPendingSplitScatter(limit) } -func (c *Controller) dispatchSplitScatterRegions() { - c.splitScatter.dispatchSplitScatterRegions() -} - func TestSplitScatterControllerCleanupResetsPendingGauge(t *testing.T) { re := require.New(t) splitScatterPendingGauge.Set(7) @@ -122,6 +118,33 @@ func TestCheckSplitScatterRegionsCreatesScatterOperator(t *testing.T) { re.Equal(group, batchGroup) } +func TestDispatchSplitScatterRegionsRequiresPrepared(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102}) + putSplitScatterRegion(tc, 101, "m", "t", splitScatterReportedCPUUsage) + putSplitScatterRegion(tc, 102, "t", "", splitScatterReportedCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + + controller.prepareChecker.ResetPrepared() + controller.dispatchSplitScatterRegions() + for _, regionID := range []uint64{100, 101, 102} { + re.Nil(oc.GetOperator(regionID)) + } + + controller.prepareChecker.SetPrepared() + controller.dispatchSplitScatterRegions() + var op *operator.Operator + for _, regionID := range []uint64{100, 101, 102} { + if op = oc.GetOperator(regionID); op != nil { + break + } + } + re.NotNil(op) +} + func TestDispatchSplitScatterKeepsPendingUntilSplitHeartbeat(t *testing.T) { re := require.New(t) controller, tc, oc, cleanup := newTestSplitScatterController(t) diff --git a/pkg/schedule/coordinator.go b/pkg/schedule/coordinator.go index bc2b3761790..75739855142 100644 --- a/pkg/schedule/coordinator.go +++ b/pkg/schedule/coordinator.go @@ -34,6 +34,7 @@ import ( "github.com/tikv/pd/pkg/schedule/diagnostic" "github.com/tikv/pd/pkg/schedule/hbstream" "github.com/tikv/pd/pkg/schedule/operator" + "github.com/tikv/pd/pkg/schedule/preparecheck" "github.com/tikv/pd/pkg/schedule/scatter" "github.com/tikv/pd/pkg/schedule/schedulers" "github.com/tikv/pd/pkg/schedule/splitter" @@ -46,9 +47,7 @@ import ( const ( runSchedulerCheckInterval = 3 * time.Second - // CollectTimeout is the timeout for collecting regions. - CollectTimeout = 5 * time.Minute - maxLoadConfigRetries = 10 + maxLoadConfigRetries = 10 // pushOperatorTickInterval is the interval try to push the operator. pushOperatorTickInterval = 500 * time.Millisecond @@ -69,7 +68,7 @@ type Coordinator struct { schedulersInitialized bool cluster sche.ClusterInformer - prepareChecker *prepareChecker + prepareChecker *preparecheck.Checker checkers *checker.Controller regionScatterer *scatter.RegionScatterer regionSplitter *splitter.RegionSplitter @@ -83,9 +82,10 @@ type Coordinator struct { // NewCoordinator creates a new Coordinator. func NewCoordinator(parentCtx context.Context, cluster sche.ClusterInformer, hbStreams *hbstream.HeartbeatStreams) *Coordinator { ctx, cancel := context.WithCancel(parentCtx) + prepareChecker := preparecheck.NewChecker(cluster.GetPrepareRegionCount) opController := operator.NewController(ctx, cluster.GetBasicCluster(), cluster.GetSharedConfig(), hbStreams) - schedulers := schedulers.NewController(ctx, cluster, cluster.GetStorage(), opController) - checkers := checker.NewController(ctx, cluster, cluster.GetCheckerConfig(), opController) + schedulers := schedulers.NewController(ctx, cluster, cluster.GetStorage(), opController, prepareChecker) + checkers := checker.NewController(ctx, cluster, cluster.GetCheckerConfig(), opController, prepareChecker) // Set the callbacks for operator success opController.SetSuccessCallbacks( @@ -101,7 +101,7 @@ func NewCoordinator(parentCtx context.Context, cluster sche.ClusterInformer, hbS cancel: cancel, schedulersInitialized: false, cluster: cluster, - prepareChecker: newPrepareChecker(cluster.GetPrepareRegionCount), + prepareChecker: prepareChecker, checkers: checkers, regionScatterer: scatter.NewRegionScatterer(ctx, cluster, opController, checkers.AddPendingProcessedRegions), regionSplitter: splitter.NewRegionSplitter(cluster, splitter.NewSplitRegionsHandler(cluster, opController), checkers.AddPendingProcessedRegions), @@ -216,6 +216,26 @@ func (c *Coordinator) driveSlowNodeScheduler() { } } +func (c *Coordinator) runPrepareChecker() { + defer logutil.LogPanic() + defer c.wg.Done() + + ticker := time.NewTicker(3 * time.Second) + defer ticker.Stop() + for { + select { + case <-c.ctx.Done(): + return + case <-ticker.C: + if !c.prepareChecker.IsPrepared() { + if c.prepareChecker.Check(c.cluster.GetBasicCluster()) { + log.Info("prepare checker is ready") + } + } + } + } +} + // RunUntilStop runs the coordinator until receiving the stop signal. func (c *Coordinator) RunUntilStop() { c.Run() @@ -249,7 +269,8 @@ func (c *Coordinator) Run() { log.Info("coordinator starts to run schedulers") c.InitSchedulers(true) - c.wg.Add(4) + c.wg.Add(5) + go c.runPrepareChecker() // Starts to patrol regions. go c.PatrolRegions() // Checks suspect key ranges @@ -581,7 +602,7 @@ func ResetHotSpotMetrics() { // ShouldRun returns true if the coordinator should run. func (c *Coordinator) ShouldRun() bool { - return c.prepareChecker.check(c.cluster.GetBasicCluster()) + return c.prepareChecker.Check(c.cluster.GetBasicCluster()) } // GetSchedulersController returns the schedulers controller. @@ -649,7 +670,7 @@ func (c *Coordinator) GetRuleChecker() *checker.RuleChecker { } // GetPrepareChecker returns the prepare checker. -func (c *Coordinator) GetPrepareChecker() *prepareChecker { +func (c *Coordinator) GetPrepareChecker() *preparecheck.Checker { return c.prepareChecker } diff --git a/pkg/schedule/prepare_checker.go b/pkg/schedule/preparecheck/prepare_checker.go similarity index 59% rename from pkg/schedule/prepare_checker.go rename to pkg/schedule/preparecheck/prepare_checker.go index 2ec4c5aaf35..a2e91dea87e 100644 --- a/pkg/schedule/prepare_checker.go +++ b/pkg/schedule/preparecheck/prepare_checker.go @@ -1,4 +1,4 @@ -// Copyright 2022 TiKV Project Authors. +// Copyright 2025 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. @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package schedule +package preparecheck import ( "time" @@ -25,28 +25,33 @@ import ( "github.com/tikv/pd/pkg/utils/syncutil" ) -type prepareChecker struct { +// collectTimeout is the timeout for collecting regions. +const collectTimeout = 5 * time.Minute + +// Checker is used to check if the coordinator has finished cluster information preparation. +type Checker struct { syncutil.RWMutex start time.Time totalRegionCountFn func() (int, error) prepared bool } -func newPrepareChecker(totalRegionCountFn func() (int, error)) *prepareChecker { - return &prepareChecker{ +// NewChecker creates a new Checker. +func NewChecker(totalRegionCountFn func() (int, error)) *Checker { + return &Checker{ start: time.Now(), totalRegionCountFn: totalRegionCountFn, } } -// Before starting up the scheduler, we need to take the proportion of the regions on each store into consideration. -func (checker *prepareChecker) check(c *core.BasicCluster) bool { +// Check is used to check if the coordinator has finished cluster information preparation. +func (checker *Checker) Check(c *core.BasicCluster) bool { checker.Lock() defer checker.Unlock() if checker.prepared { return true } - if time.Since(checker.start) > CollectTimeout { + if time.Since(checker.start) > collectTimeout { checker.prepared = true return true } @@ -78,7 +83,7 @@ func (checker *prepareChecker) check(c *core.BasicCluster) bool { } // IsPrepared returns whether the coordinator is prepared. -func (checker *prepareChecker) IsPrepared() bool { +func (checker *Checker) IsPrepared() bool { if checker == nil { return false } @@ -87,9 +92,42 @@ func (checker *prepareChecker) IsPrepared() bool { return checker.prepared } +// RunIfPrepared runs fn while the checker remains prepared. +func (checker *Checker) RunIfPrepared(fn func()) bool { + if checker == nil { + return false + } + checker.RLock() + defer checker.RUnlock() + if !checker.prepared { + return false + } + fn() + return true +} + // SetPrepared is for test purpose -func (checker *prepareChecker) SetPrepared() { +func (checker *Checker) SetPrepared() { checker.Lock() defer checker.Unlock() checker.prepared = true } + +// ResetPrepared resets the checker so scheduling pauses until cluster +// information is collected again after the region cache is reset. +func (checker *Checker) ResetPrepared() { + checker.ResetPreparedAndRun(nil) +} + +// ResetPreparedAndRun resets the checker and runs fn while scheduling is +// fenced. It is used to keep the prepared state and the region cache reset +// atomic with respect to prepare checks and scheduling. +func (checker *Checker) ResetPreparedAndRun(fn func()) { + checker.Lock() + defer checker.Unlock() + checker.prepared = false + checker.start = time.Now() + if fn != nil { + fn() + } +} diff --git a/pkg/schedule/preparecheck/prepare_checker_test.go b/pkg/schedule/preparecheck/prepare_checker_test.go new file mode 100644 index 00000000000..9a6893ffbac --- /dev/null +++ b/pkg/schedule/preparecheck/prepare_checker_test.go @@ -0,0 +1,107 @@ +// 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 preparecheck + +import ( + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.uber.org/goleak" + + "github.com/tikv/pd/pkg/core" + "github.com/tikv/pd/pkg/utils/testutil" +) + +func TestMain(m *testing.M) { + goleak.VerifyTestMain(m, testutil.LeakOptions...) +} + +func TestRunIfPrepared(t *testing.T) { + re := require.New(t) + checker := NewChecker(func() (int, error) { + return 0, nil + }) + + ran := false + re.False(checker.RunIfPrepared(func() { + ran = true + })) + re.False(ran) + + checker.SetPrepared() + re.True(checker.RunIfPrepared(func() { + ran = true + })) + re.True(ran) + + ran = false + checker.ResetPrepared() + re.False(checker.RunIfPrepared(func() { + ran = true + })) + re.False(ran) +} + +func TestResetPreparedAndRunFencesPrepareCheck(t *testing.T) { + re := require.New(t) + var resetFinished atomic.Bool + var checkedBeforeResetFinished atomic.Bool + checker := NewChecker(func() (int, error) { + if !resetFinished.Load() { + checkedBeforeResetFinished.Store(true) + } + return 0, nil + }) + checker.SetPrepared() + + resetStarted := make(chan struct{}) + finishReset := make(chan struct{}) + resetDone := make(chan struct{}) + go func() { + checker.ResetPreparedAndRun(func() { + close(resetStarted) + <-finishReset + resetFinished.Store(true) + }) + close(resetDone) + }() + <-resetStarted + + checkStarted := make(chan struct{}) + checkDone := make(chan bool, 1) + go func() { + close(checkStarted) + checkDone <- checker.Check(core.NewBasicCluster()) + }() + <-checkStarted + time.Sleep(50 * time.Millisecond) + completedEarly := false + select { + case <-checkDone: + completedEarly = true + default: + } + + close(finishReset) + <-resetDone + if completedEarly { + re.Fail("prepare check completed before reset callback") + return + } + re.True(<-checkDone) + re.False(checkedBeforeResetFinished.Load()) +} diff --git a/pkg/schedule/schedulers/scheduler_controller.go b/pkg/schedule/schedulers/scheduler_controller.go index c5fec568685..ceaa0001345 100644 --- a/pkg/schedule/schedulers/scheduler_controller.go +++ b/pkg/schedule/schedulers/scheduler_controller.go @@ -31,6 +31,7 @@ import ( "github.com/tikv/pd/pkg/schedule/labeler" "github.com/tikv/pd/pkg/schedule/operator" "github.com/tikv/pd/pkg/schedule/plan" + "github.com/tikv/pd/pkg/schedule/preparecheck" "github.com/tikv/pd/pkg/schedule/types" "github.com/tikv/pd/pkg/storage/endpoint" "github.com/tikv/pd/pkg/utils/logutil" @@ -57,10 +58,17 @@ type Controller struct { // which will only be initialized and used in the microservice env now. schedulerHandlers map[string]http.Handler opController *operator.Controller + prepareChecker *preparecheck.Checker } // NewController creates a scheduler controller. -func NewController(ctx context.Context, cluster sche.SchedulerCluster, storage endpoint.ConfigStorage, opController *operator.Controller) *Controller { +func NewController(ctx context.Context, cluster sche.SchedulerCluster, storage endpoint.ConfigStorage, opController *operator.Controller, prepareCheckers ...*preparecheck.Checker) *Controller { + prepareChecker := preparecheck.NewChecker(cluster.GetPrepareRegionCount) + if len(prepareCheckers) > 0 && prepareCheckers[0] != nil { + prepareChecker = prepareCheckers[0] + } else { + prepareChecker.SetPrepared() + } return &Controller{ ctx: ctx, cluster: cluster, @@ -68,6 +76,7 @@ func NewController(ctx context.Context, cluster sche.SchedulerCluster, storage e schedulers: make(map[string]*ScheduleController), schedulerHandlers: make(map[string]http.Handler), opController: opController, + prepareChecker: prepareChecker, } } @@ -389,13 +398,22 @@ func (c *Controller) runScheduler(s *ScheduleController) { for { select { case <-ticker.C: - diagnosable := s.IsDiagnosticAllowed() - if !s.AllowSchedule(diagnosable) { + shouldResetTicker := false + if !c.prepareChecker.RunIfPrepared(func() { + diagnosable := s.IsDiagnosticAllowed() + if !s.AllowSchedule(diagnosable) { + return + } + shouldResetTicker = true + if op := s.Schedule(diagnosable); len(op) > 0 { + added := c.opController.AddWaitingOperator(op...) + log.Debug("add operator", zap.Int("added", added), zap.Int("total", len(op)), zap.String("scheduler", s.GetName())) + } + }) { continue } - if op := s.Schedule(diagnosable); len(op) > 0 { - added := c.opController.AddWaitingOperator(op...) - log.Debug("add operator", zap.Int("added", added), zap.Int("total", len(op)), zap.String("scheduler", s.GetName())) + if !shouldResetTicker { + continue } // Note: we reset the ticker here to support updating configuration dynamically. ticker.Reset(s.GetInterval()) diff --git a/pkg/unsaferecovery/unsafe_recovery_controller.go b/pkg/unsaferecovery/unsafe_recovery_controller.go index d0a27ab9871..8498402954c 100644 --- a/pkg/unsaferecovery/unsafe_recovery_controller.go +++ b/pkg/unsaferecovery/unsafe_recovery_controller.go @@ -16,6 +16,7 @@ package unsaferecovery import ( "bytes" + "context" "encoding/json" "fmt" "sort" @@ -46,6 +47,7 @@ type stage int const ( defaultPlanExecutionTimeout = time.Second * 60 + resetRegionCacheTimeout = 3 * time.Second ) var globalRecoveryStep = uint64(time.Now().UnixNano()) @@ -94,12 +96,12 @@ func nextRecoveryStep() uint64 { // | | // | | // v | -// +-----------+ +-----------+ | -// +-----------+ | | | | | -// | | | ExitForce | | Create | | -// | Finished |<------| Leader |<----------| Region |-----+ -// | | | | | | -// +-----------+ +-----------+ +-----------+ +// +-----------+ +-----------+ +-----------+ | +// +-----------+ | | | | | | | +// | | | Resetting | | ExitForce | | Create | | +// | Finished |<------| Cache |<------| Leader |<------| Region |-----+ +// | | | | | | | | +// +-----------+ +-----------+ +-----------+ +-----------+ const ( Idle stage = iota CollectReport @@ -109,6 +111,7 @@ const ( DemoteFailedVoter CreateEmptyRegion ExitForceLeader + Resetting Finished Failed ) @@ -116,7 +119,8 @@ const ( type cluster interface { core.StoreSetInformer - ResetRegionCache() + Context() context.Context + ResetPreparedAndResetRegionCache(context.Context) error AllocID(uint32) (uint64, uint32, error) BuryStore(storeID uint64, forceBury bool) error GetSchedulerConfig() sc.SchedulerConfigProvider @@ -157,10 +161,12 @@ type Controller struct { // accumulated output for the whole recovery process output []StageOutput // exposed to the outside for testing - AffectedTableIDs map[int64]struct{} - affectedMetaRegions map[uint64]struct{} - newlyCreatedRegions map[uint64]struct{} - err error + AffectedTableIDs map[int64]struct{} + affectedMetaRegions map[uint64]struct{} + newlyCreatedRegions map[uint64]struct{} + resettingRegionCache bool + resetRegionCacheErr error + err error } // StageOutput is the information for one stage of the recovery process. @@ -194,6 +200,8 @@ func (u *Controller) reset() { u.affectedMetaRegions = make(map[uint64]struct{}, 0) u.newlyCreatedRegions = make(map[uint64]struct{}, 0) u.orphanedPeers = make(map[uint64][]*metapb.Peer) + u.resettingRegionCache = false + u.resetRegionCacheErr = nil u.err = nil u.planExecutionTimeout = defaultPlanExecutionTimeout u.disableParanoidCheck = false @@ -292,6 +300,9 @@ func (u *Controller) AbortFailedStoresRemoval() error { if !isRunning(u.stage) { return errs.ErrUnsafeRecoveryInvalidInput.FastGenByArgs("no ongoing unsafe recovery") } + if u.stage == Resetting { + return errs.ErrUnsafeRecoveryInvalidInput.FastGenByArgs("unsafe recovery is resetting region cache") + } if u.stage == ExitForceLeader { return nil } @@ -321,6 +332,13 @@ func (u *Controller) Show() []StageOutput { func (u *Controller) getReportStatus() StageOutput { var status StageOutput status.Time = time.Now().Format("2006-01-02 15:04:05.000") + if u.stage == Resetting { + status.Info = "Resetting region cache before finishing unsafe recovery" + if u.resetRegionCacheErr != nil { + status.Details = append(status.Details, fmt.Sprintf("failed to reset region cache: %v", u.resetRegionCacheErr)) + } + return status + } if u.numStoresReported != len(u.storeReports) { status.Info = fmt.Sprintf("Collecting reports from alive stores(%d/%d)", u.numStoresReported, len(u.storeReports)) var ( @@ -350,7 +368,7 @@ func (u *Controller) getReportStatus() StageOutput { } func (u *Controller) checkTimeout() error { - if u.stage == Finished || u.stage == Failed { + if u.stage == Resetting || u.stage == Finished || u.stage == Failed { return nil } @@ -388,10 +406,18 @@ func (u *Controller) handleErr(err error) bool { // send detailed report back. func (u *Controller) HandleStoreHeartbeat(heartbeat *pdpb.StoreHeartbeatRequest, resp *pdpb.StoreHeartbeatResponse) { u.Lock() - defer u.Unlock() if !isRunning(u.stage) { // no recovery in progress, do nothing + u.Unlock() + return + } + if u.stage == Resetting { + shouldReset := u.startRegionCacheReset() + u.Unlock() + if shouldReset { + u.resetRegionCacheAndFinish() + } return } @@ -417,9 +443,48 @@ func (u *Controller) HandleStoreHeartbeat(heartbeat *pdpb.StoreHeartbeatRequest, }() if done || (err != nil && u.handleErr(err)) { + shouldReset := u.startRegionCacheReset() + u.Unlock() + if shouldReset { + u.resetRegionCacheAndFinish() + } return } u.dispatchPlan(heartbeat, resp) + u.Unlock() +} + +// startRegionCacheReset marks the cache reset as in progress. The caller must +// hold the controller lock. +func (u *Controller) startRegionCacheReset() bool { + if u.stage != Resetting || u.resettingRegionCache { + return false + } + u.resettingRegionCache = true + return true +} + +func (u *Controller) resetRegionCacheAndFinish() { + parentCtx := u.cluster.Context() + if parentCtx == nil { + parentCtx = context.Background() + } + ctx, cancel := context.WithTimeout(parentCtx, resetRegionCacheTimeout) + err := u.cluster.ResetPreparedAndResetRegionCache(ctx) + cancel() + + u.Lock() + defer u.Unlock() + u.resettingRegionCache = false + if u.stage != Resetting { + return + } + u.resetRegionCacheErr = err + if err != nil { + log.Warn("failed to reset region cache before finishing unsafe recovery", zap.Error(err)) + return + } + u.changeStage(Finished) } func (u *Controller) generatePlan(newestRegionTree *regionTree, peersMap map[uint64][]*regionItem) (bool, error) { @@ -494,6 +559,10 @@ func (u *Controller) generatePlan(newestRegionTree *regionTree, peersMap map[uin if err == nil && !hasPlan { if u.err != nil { u.changeStage(Failed) + } else if u.step > u.recoveryStartStep+1 { + // Only CollectReport has finished when step == recoveryStartStep+1, + // which means no operation has done and no cache invalidation is needed. + u.changeStage(Resetting) } else { u.changeStage(Finished) } @@ -612,12 +681,9 @@ func (u *Controller) changeStage(stage stage) { if u.err != nil { output.Details = append(output.Details, fmt.Sprintf("triggered by error: %v", u.err.Error())) } + case Resetting: + output.Info = "Unsafe recovery is resetting region cache" case Finished: - if u.step > u.recoveryStartStep+1 { - // Only CollectReport has finished when step == recoveryStartStep+1, - // which means no operation has done and no cache invalidation is needed. - u.cluster.ResetRegionCache() - } output.Info = "Unsafe recovery Finished" output.Details = u.getAffectedTableDigest() u.storePlanExpires = make(map[uint64]time.Time) diff --git a/pkg/unsaferecovery/unsafe_recovery_controller_test.go b/pkg/unsaferecovery/unsafe_recovery_controller_test.go index ec1ee7eb1ff..9ebbe74bb3c 100644 --- a/pkg/unsaferecovery/unsafe_recovery_controller_test.go +++ b/pkg/unsaferecovery/unsafe_recovery_controller_test.go @@ -17,6 +17,7 @@ package unsaferecovery import ( "context" + "errors" "fmt" "testing" "time" @@ -42,6 +43,23 @@ func TestMain(m *testing.M) { goleak.VerifyTestMain(m, testutil.LeakOptions...) } +type resetTrackingCluster struct { + *mockcluster.Cluster + resetCalls int + resetErr error + resetContext context.Context + onReset func() +} + +func (c *resetTrackingCluster) ResetPreparedAndResetRegionCache(ctx context.Context) error { + c.resetCalls++ + c.resetContext = ctx + if c.onReset != nil { + c.onReset() + } + return c.resetErr +} + func newStoreHeartbeat(storeID uint64, report *pdpb.StoreReport) *pdpb.StoreHeartbeatRequest { return &pdpb.StoreHeartbeatRequest{ Stats: &pdpb.StoreStats{ @@ -2036,6 +2054,102 @@ func TestUnsafeRecoveryStepIsUniqueAcrossRuns(t *testing.T) { re.Nil(recoveryController.storeReports[1]) } +func TestFinishRetriesCacheResetBeforeLeavingRunningStage(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cluster := &resetTrackingCluster{ + Cluster: mockcluster.NewCluster(ctx, mockconfig.NewTestOptions()), + resetErr: errors.New("scheduling cache reset failed"), + } + recoveryController := NewController(cluster) + recoveryController.stage = ExitForceLeader + recoveryController.recoveryStartStep = 1 + recoveryController.step = 3 + recoveryController.timeout = time.Now().Add(time.Minute) + recoveryController.storeReports[1] = nil + + recoveryController.HandleStoreHeartbeat( + newStoreHeartbeat(1, &pdpb.StoreReport{Step: recoveryController.step}), + &pdpb.StoreHeartbeatResponse{}, + ) + + re.Equal(1, cluster.resetCalls) + _, hasDeadline := cluster.resetContext.Deadline() + re.True(hasDeadline) + re.True(recoveryController.IsRunning()) + re.Equal(Resetting, recoveryController.GetStage()) + + recoveryController.timeout = time.Now().Add(-time.Second) + status := recoveryController.Show() + re.Equal(Resetting, recoveryController.GetStage()) + re.Contains(status[len(status)-1].Info, "Resetting region cache") + re.Contains(status[len(status)-1].Details[0], "scheduling cache reset failed") + re.ErrorContains(recoveryController.AbortFailedStoresRemoval(), "resetting region cache") + + cluster.resetErr = nil + recoveryController.HandleStoreHeartbeat( + newStoreHeartbeat(1, &pdpb.StoreReport{}), + &pdpb.StoreHeartbeatResponse{}, + ) + re.Equal(2, cluster.resetCalls) + re.Equal(Finished, recoveryController.GetStage()) +} + +func TestFinishDoesNotHoldControllerLockDuringCacheReset(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + resetStarted := make(chan struct{}) + finishReset := make(chan struct{}) + cluster := &resetTrackingCluster{ + Cluster: mockcluster.NewCluster(ctx, mockconfig.NewTestOptions()), + onReset: func() { + close(resetStarted) + <-finishReset + }, + } + recoveryController := NewController(cluster) + recoveryController.stage = ExitForceLeader + recoveryController.recoveryStartStep = 1 + recoveryController.step = 3 + recoveryController.timeout = time.Now().Add(time.Minute) + recoveryController.storeReports[1] = nil + + handleDone := make(chan struct{}) + go func() { + recoveryController.HandleStoreHeartbeat( + newStoreHeartbeat(1, &pdpb.StoreReport{Step: recoveryController.step}), + &pdpb.StoreHeartbeatResponse{}, + ) + close(handleDone) + }() + <-resetStarted + + isRunningDone := make(chan bool, 1) + go func() { + isRunningDone <- recoveryController.IsRunning() + }() + lockWasAvailable := false + select { + case running := <-isRunningDone: + lockWasAvailable = running + case <-time.After(100 * time.Millisecond): + } + + close(finishReset) + <-handleDone + if !lockWasAvailable { + select { + case <-isRunningDone: + default: + } + } + re.True(lockWasAvailable) +} + func TestRunning(t *testing.T) { re := require.New(t) ctx, cancel := context.WithCancel(context.Background()) diff --git a/server/api/admin.go b/server/api/admin.go index f298ee4cb0c..567f657da28 100644 --- a/server/api/admin.go +++ b/server/api/admin.go @@ -81,7 +81,7 @@ func (h *adminHandler) DeleteRegionCache(w http.ResponseWriter, r *http.Request) rc.RemoveRegionIfExist(regionID) msg := "The region is removed from server cache." if rc.IsServiceIndependent(constant.SchedulingServiceName) { - err = h.deleteRegionCacheInSchedulingServer(regionID) + err = h.svr.DeleteRegionCacheInSchedulingServer(r.Context(), regionID) if err != nil { msg = buildMsg(err) } @@ -126,7 +126,7 @@ func (h *adminHandler) DeleteRegionStorage(w http.ResponseWriter, r *http.Reques rc.RemoveRegionIfExist(regionID) msg := "The region is removed from server cache and region meta storage." if rc.IsServiceIndependent(constant.SchedulingServiceName) { - err = h.deleteRegionCacheInSchedulingServer(regionID) + err = h.svr.DeleteRegionCacheInSchedulingServer(r.Context(), regionID) if err != nil { msg = buildMsg(err) } @@ -154,13 +154,10 @@ func (h *adminHandler) DeleteAllRegionCache(w http.ResponseWriter, r *http.Reque h.rd.JSON(w, http.StatusOK, "All regions are removed from follower cache and the follower starts to resync regions from leader.") return } - rc.ResetRegionCache() + err = rc.ResetPreparedAndResetRegionCache(r.Context()) msg := "All regions are removed from server cache." - if rc.IsServiceIndependent(constant.SchedulingServiceName) { - err = h.deleteRegionCacheInSchedulingServer() - if err != nil { - msg = buildMsg(err) - } + if err != nil { + msg = buildMsg(err) } h.rd.JSON(w, http.StatusOK, msg) @@ -280,27 +277,6 @@ func (h *adminHandler) recoverAllocID(w http.ResponseWriter, r *http.Request) { h.rd.Text(w, http.StatusOK, "") } -func (h *adminHandler) deleteRegionCacheInSchedulingServer(id ...uint64) error { - addr, ok := h.svr.GetServicePrimaryAddr(h.svr.Context(), constant.SchedulingServiceName) - if !ok { - return errs.ErrNotFoundSchedulingPrimary.FastGenByArgs() - } - var idStr string - if len(id) > 0 { - idStr = strconv.FormatUint(id[0], 10) - } - url := fmt.Sprintf("%s/scheduling/api/v1/admin/cache/regions/%s", addr, idStr) - resp, err := apiutil.DoDelete(h.svr.GetHTTPClient(), url) - if err != nil { - return err - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return errs.ErrSchedulingServer.FastGenByArgs(resp.StatusCode) - } - return nil -} - func buildMsg(err error) string { return fmt.Sprintf("This operation was executed in PD but needs to be re-executed on scheduling server due to the following error: %s", err.Error()) } diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index 5cdb02cb23f..09a247329a4 100644 --- a/server/cluster/cluster.go +++ b/server/cluster/cluster.go @@ -206,7 +206,8 @@ type RaftCluster struct { // syncRegionRunner is used to sync region asynchronously. syncRegionRunner ratelimit.Runner - stopGCStateManager func() + stopGCStateManager func() + resetSchedulingCacheFunc func(context.Context) error } // Status saves some state information. @@ -2673,6 +2674,26 @@ func (c *RaftCluster) IsServiceIndependent(name string) bool { return exist } +// SetResetSchedulingCacheFunc sets the callback used to reset the cache in an +// independent scheduling service. +func (c *RaftCluster) SetResetSchedulingCacheFunc(fn func(context.Context) error) { + c.resetSchedulingCacheFunc = fn +} + +// ResetPreparedAndResetRegionCache atomically pauses local scheduling and +// clears the local region cache. It also resets the active scheduling service +// when scheduling runs independently. +func (c *RaftCluster) ResetPreparedAndResetRegionCache(ctx context.Context) error { + c.resetPreparedAndRegionCache() + if !c.IsServiceIndependent(constant.SchedulingServiceName) { + return nil + } + if c.resetSchedulingCacheFunc == nil { + return errs.ErrNotFoundSchedulingPrimary.FastGenByArgs() + } + return c.resetSchedulingCacheFunc(ctx) +} + // SetServiceIndependent sets the service to be independent. func (c *RaftCluster) SetServiceIndependent(name string) { c.independentServices.Store(name, struct{}{}) diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go index b8080d7a45f..c675d278536 100644 --- a/server/cluster/cluster_test.go +++ b/server/cluster/cluster_test.go @@ -3061,6 +3061,7 @@ func TestCheckCache(t *testing.T) { cfg.ReplicaScheduleLimit = 0 }, nil, nil, re) defer cleanup() + co.GetPrepareChecker().SetPrepared() oc := co.GetOperatorController() checker := co.GetCheckerController() @@ -3156,6 +3157,7 @@ func TestPatrolRegionConcurrency(t *testing.T) { cfg.MergeScheduleLimit = uint64(mergeScheduleLimit) }, nil, nil, re) defer cleanup() + co.GetPrepareChecker().SetPrepared() oc := co.GetOperatorController() checker := co.GetCheckerController() @@ -3218,6 +3220,8 @@ func TestScanLimit(t *testing.T) { func checkScanLimit(re *require.Assertions, regionCount int, expectScanLimit ...int) { tc, co, cleanup := prepare(nil, nil, nil, re) defer cleanup() + // set prepared to avoid prepare checker block the patrol + co.GetPrepareChecker().SetPrepared() re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/schedule/checker/breakPatrol", `return`)) re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/schedule/checker/regionCount", fmt.Sprintf("return(\"%d\")", regionCount))) defer func() { diff --git a/server/cluster/scheduling_controller.go b/server/cluster/scheduling_controller.go index f9819e8dc53..a69d553f286 100644 --- a/server/cluster/scheduling_controller.go +++ b/server/cluster/scheduling_controller.go @@ -484,6 +484,19 @@ func (sc *schedulingController) SetPrepared() { sc.coordinator.GetPrepareChecker().SetPrepared() } +// ResetPrepared reset the prepare checker. +func (sc *schedulingController) ResetPrepared() { + sc.mu.RLock() + defer sc.mu.RUnlock() + sc.coordinator.GetPrepareChecker().ResetPrepared() +} + +func (sc *schedulingController) resetPreparedAndRegionCache() { + sc.mu.RLock() + defer sc.mu.RUnlock() + sc.coordinator.GetPrepareChecker().ResetPreparedAndRun(sc.ResetRegionCache) +} + // IsSchedulingControllerRunning returns whether the scheduling controller is running. Only for test purpose. func (sc *schedulingController) IsSchedulingControllerRunning() bool { sc.mu.RLock() diff --git a/server/cluster/split_scatter_test.go b/server/cluster/split_scatter_test.go index a6e0793631f..701dd267fc3 100644 --- a/server/cluster/split_scatter_test.go +++ b/server/cluster/split_scatter_test.go @@ -28,6 +28,7 @@ import ( "github.com/tikv/pd/pkg/codec" "github.com/tikv/pd/pkg/core" + "github.com/tikv/pd/pkg/mcs/utils/constant" "github.com/tikv/pd/pkg/mock/mockid" "github.com/tikv/pd/pkg/schedule/hbstream" "github.com/tikv/pd/pkg/schedule/labeler" @@ -43,6 +44,23 @@ const ( splitScatterReportedCPUUsage uint64 = 1 ) +func TestResetPreparedAndResetRegionCacheForwardsToScheduling(t *testing.T) { + re := require.New(t) + cluster, _ := newSplitScatterTestCluster(t) + cluster.SetServiceIndependent(constant.SchedulingServiceName) + + forwarded := 0 + cluster.SetResetSchedulingCacheFunc(func(context.Context) error { + forwarded++ + return nil + }) + re.NoError(cluster.ResetPreparedAndResetRegionCache(context.Background())) + + re.Equal(1, forwarded) + re.False(cluster.GetCoordinator().GetPrepareChecker().IsPrepared()) + re.Zero(cluster.GetTotalRegionCount()) +} + func TestHandleAskBatchSplitSchedulesSplitScatterInPatrol(t *testing.T) { re := require.New(t) cluster, cancelPatrol := newSplitScatterTestCluster(t) @@ -168,6 +186,7 @@ func newSplitScatterTestCluster(t *testing.T) (*RaftCluster, context.CancelFunc) re.NoError(err) hbStreams := hbstream.NewTestHeartbeatStreams(ctx, cluster.BasicCluster, false) cluster.initCoordinator(ctx, cluster, hbStreams) + cluster.GetCoordinator().GetPrepareChecker().SetPrepared() t.Cleanup(func() { hbStreams.Close() }) diff --git a/server/server.go b/server/server.go index 58e3caa31e6..d620896f035 100644 --- a/server/server.go +++ b/server/server.go @@ -525,6 +525,9 @@ func (s *Server) startServer(ctx context.Context) error { s.tsoAllocator = tso.NewAllocator(s.ctx, constant.DefaultKeyspaceGroupID, s.member, tsoStorage, s) s.basicCluster = core.NewBasicCluster() s.cluster = cluster.NewRaftCluster(ctx, s.GetMember(), s.GetBasicCluster(), s.GetStorage(), syncer.NewRegionSyncer(s), s.client, s.httpClient, s.tsoAllocator) + s.cluster.SetResetSchedulingCacheFunc(func(ctx context.Context) error { + return s.DeleteRegionCacheInSchedulingServer(ctx) + }) keyspaceIDAllocator := id.NewAllocator(&id.AllocatorParams{ Client: s.client, Label: id.KeyspaceLabel, @@ -2332,6 +2335,29 @@ func (s *Server) GetServicePrimaryAddr(ctx context.Context, serviceName string) return "", false } +// DeleteRegionCacheInSchedulingServer deletes all or one region from the cache +// of the primary scheduling server. +func (s *Server) DeleteRegionCacheInSchedulingServer(ctx context.Context, id ...uint64) error { + addr, ok := s.GetServicePrimaryAddr(ctx, mcs.SchedulingServiceName) + if !ok { + return errs.ErrNotFoundSchedulingPrimary.FastGenByArgs() + } + var idStr string + if len(id) > 0 { + idStr = strconv.FormatUint(id[0], 10) + } + url := fmt.Sprintf("%s/scheduling/api/v1/admin/cache/regions/%s", addr, idStr) + resp, err := apiutil.DoDeleteWithContext(ctx, s.GetHTTPClient(), url) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return errs.ErrSchedulingServer.FastGenByArgs(resp.StatusCode) + } + return nil +} + // SetServicePrimaryAddr sets the primary address directly. // Note: This function is only used for test. func (s *Server) SetServicePrimaryAddr(serviceName, addr string) { diff --git a/server/server_test.go b/server/server_test.go index 848104355f2..155fdf777a6 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -17,6 +17,8 @@ package server import ( "context" stderrors "errors" + "net/http" + "net/http/httptest" "testing" "github.com/stretchr/testify/require" @@ -24,6 +26,7 @@ import ( "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" + "github.com/tikv/pd/pkg/mcs/utils/constant" "github.com/tikv/pd/pkg/member" "github.com/tikv/pd/pkg/storage" "github.com/tikv/pd/pkg/storage/kv" @@ -33,6 +36,33 @@ import ( var errTestFollowerRegionStorage = stderrors.New("test follower region storage error") +func TestDeleteRegionCacheInSchedulingServer(t *testing.T) { + re := require.New(t) + requests := make(chan *http.Request, 2) + schedulingServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests <- r.Clone(r.Context()) + w.WriteHeader(http.StatusOK) + })) + defer schedulingServer.Close() + + s := &Server{ctx: context.Background(), httpClient: schedulingServer.Client()} + s.SetServicePrimaryAddr(constant.SchedulingServiceName, schedulingServer.URL) + + re.NoError(s.DeleteRegionCacheInSchedulingServer(context.Background())) + request := <-requests + re.Equal(http.MethodDelete, request.Method) + re.Equal("/scheduling/api/v1/admin/cache/regions/", request.URL.Path) + + re.NoError(s.DeleteRegionCacheInSchedulingServer(context.Background(), 42)) + request = <-requests + re.Equal(http.MethodDelete, request.Method) + re.Equal("/scheduling/api/v1/admin/cache/regions/42", request.URL.Path) + + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + re.ErrorIs(s.DeleteRegionCacheInSchedulingServer(canceledCtx), context.Canceled) +} + func TestResetFollowerRegionCacheRequiresRegionStorage(t *testing.T) { re := require.New(t) cfg := config.NewConfig() diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go index d13f442d736..9a8964e059e 100644 --- a/tests/server/api/api_test.go +++ b/tests/server/api/api_test.go @@ -38,6 +38,7 @@ import ( "github.com/tikv/pd/pkg/core" "github.com/tikv/pd/pkg/response" + "github.com/tikv/pd/pkg/schedule/operator" "github.com/tikv/pd/pkg/utils/apiutil" "github.com/tikv/pd/pkg/utils/testutil" "github.com/tikv/pd/pkg/utils/typeutil" @@ -172,7 +173,6 @@ func (suite *middlewareTestSuite) TestRequestInfoMiddleware() { resp.Body.Close() re.NoError(err) re.Equal(http.StatusOK, resp.StatusCode) - re.Equal("Profile", resp.Header.Get("service-label")) re.JSONEq("{\"seconds\":[\"1\"]}", resp.Header.Get("url-param")) re.JSONEq("{\"testkey\":\"testvalue\"}", resp.Header.Get("body-param")) @@ -1694,3 +1694,81 @@ func sendRequest(re *require.Assertions, url string, method string, statusCode i return output } + +func TestDeleteAllRegionCacheScheduling(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cluster, err := tests.NewTestCluster(ctx, 1) + re.NoError(err) + defer cluster.Destroy() + + // Initialize the cluster + re.NoError(cluster.RunInitialServers()) + re.NotEmpty(cluster.WaitLeader()) + leaderServer := cluster.GetLeaderServer() + re.NotNil(leaderServer) + re.NoError(leaderServer.BootstrapCluster()) + rc := leaderServer.GetRaftCluster() + re.NotNil(rc) + + // Add 3 storage nodes + for i := uint64(1); i <= 3; i++ { + store := &metapb.Store{ + Id: i, + Address: fmt.Sprintf("tikv%d", i), + State: metapb.StoreState_Up, + NodeState: metapb.NodeState_Serving, + LastHeartbeat: time.Now().UnixNano(), + } + tests.MustPutStore(re, cluster, store) + } + + // Create a test region + region := &metapb.Region{ + Id: 2, + StartKey: []byte(""), + EndKey: []byte(""), + RegionEpoch: &metapb.RegionEpoch{ + ConfVer: 1, + Version: 1, + }, + Peers: []*metapb.Peer{ + {Id: 11, StoreId: 1, Role: metapb.PeerRole_Voter}, + {Id: 12, StoreId: 2, Role: metapb.PeerRole_Voter}, + }, + } + regionInfo := core.NewRegionInfo(region, region.Peers[0], core.SetSource(core.Heartbeat)) + + // Wait for the cluster to be ready + testutil.Eventually(re, func() bool { + return rc.GetCoordinator() != nil + }) + err = rc.HandleRegionHeartbeat(regionInfo) + re.NoError(err) + testutil.Eventually(re, func() bool { + return rc.GetCoordinator().GetPrepareChecker().IsPrepared() + }) + + testutil.Eventually(re, func() bool { + count := rc.GetOperatorController().OperatorCount(operator.OpSplit) + return count > 0 + }) + + // Call the delete API and verify + addr := leaderServer.GetAddr() + "/pd/api/v1/admin/cache/regions" + output := sendRequest(re, addr, http.MethodDelete, http.StatusOK) + re.Contains(string(output), "All regions are removed from server cache") + re.False(rc.GetCoordinator().GetPrepareChecker().IsPrepared()) + rc.GetOperatorController().RemoveOperators() + re.Equal(0, int(rc.GetOperatorController().OperatorCount(operator.OpSplit))) + + // Simulate continuous heartbeat and verify scheduling recovery + err = rc.HandleRegionHeartbeat(regionInfo) + re.NoError(err) + + testutil.Eventually(re, func() bool { + count := rc.GetOperatorController().OperatorCount(operator.OpSplit) + return count > 0 + }) +} diff --git a/tests/server/api/diagnostic_test.go b/tests/server/api/diagnostic_test.go index 5572bb2d726..eab06c5aa2d 100644 --- a/tests/server/api/diagnostic_test.go +++ b/tests/server/api/diagnostic_test.go @@ -68,6 +68,8 @@ func (suite *diagnosticTestSuite) checkSchedulerDiagnosticAPI(cluster *tests.Tes } leader := cluster.GetLeaderServer() + // Set prepared to allow schedulers to run and generate diagnostic results + leader.GetRaftCluster().GetCoordinator().GetPrepareChecker().SetPrepared() urlPrefix := leader.GetAddr() + "/pd/api/v1" diagnosticPrefix := fmt.Sprintf("%s/schedulers/diagnostic", urlPrefix) diff --git a/tests/server/api/region_test.go b/tests/server/api/region_test.go index 550437ca2a5..757424c0abc 100644 --- a/tests/server/api/region_test.go +++ b/tests/server/api/region_test.go @@ -55,10 +55,14 @@ func TestRegionTestSuite(t *testing.T) { } func (suite *regionTestSuite) SetupSuite() { + re := suite.Require() + re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/schedule/checker/skipCheckSuspectRanges", "return(true)")) suite.env = tests.NewSchedulingTestEnvironment(suite.T()) } func (suite *regionTestSuite) TearDownSuite() { + re := suite.Require() + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/schedule/checker/skipCheckSuspectRanges")) suite.env.Cleanup() } diff --git a/tests/server/cluster/cluster_test.go b/tests/server/cluster/cluster_test.go index b66785628ab..9d117d672fe 100644 --- a/tests/server/cluster/cluster_test.go +++ b/tests/server/cluster/cluster_test.go @@ -1947,9 +1947,12 @@ func TestExternalTimestamp(t *testing.T) { re := require.New(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() + tc, err := tests.NewTestCluster(ctx, 1) defer tc.Destroy() re.NoError(err) + fname := testutil.InitTempFileLogger("debug") + defer os.RemoveAll(fname) err = tc.RunInitialServers() re.NoError(err) tc.WaitLeader() diff --git a/tests/server/storage/hot_region_storage_test.go b/tests/server/storage/hot_region_storage_test.go index 45625a6f980..4ea41e69c21 100644 --- a/tests/server/storage/hot_region_storage_test.go +++ b/tests/server/storage/hot_region_storage_test.go @@ -151,10 +151,11 @@ func (s *hotRegionStorageTestSuite) checkHotRegionStorage(cluster *tests.TestClu next, err = iter.Next() re.NoError(err) re.Nil(next) - iter = hotRegionStorage.NewIterator([]string{utils.Read.String()}, startTime*1000, time.Now().UnixMilli()) - next, err = iter.Next() - re.NoError(err) - re.NotNil(next) + testutil.Eventually(re, func() bool { // wait for the read history hot region to be written to the storage + iter = hotRegionStorage.NewIterator([]string{utils.Read.String()}, startTime*1000, time.Now().UnixMilli()) + next, err = iter.Next() + return err == nil && next != nil + }) re.Equal(uint64(3), next.RegionID) re.Equal(uint64(1), next.StoreID) re.Equal(utils.Read.String(), next.HotRegionType) diff --git a/tools/pd-ctl/pdctl/command/operator_command.go b/tools/pd-ctl/pdctl/command/operator_command.go index 8e3fcc3ead6..3670f63a64f 100644 --- a/tools/pd-ctl/pdctl/command/operator_command.go +++ b/tools/pd-ctl/pdctl/command/operator_command.go @@ -274,13 +274,13 @@ func NewAddLearnerCommand() *cobra.Command { func addLearnerCommandFunc(cmd *cobra.Command, args []string) { if len(args) != 2 { - fmt.Println(cmd.UsageString()) + cmd.Println(cmd.UsageString()) return } ids, err := parseUint64s(args) if err != nil { - fmt.Println(err) + cmd.Println(err) return } diff --git a/tools/pd-ctl/tests/store/store_test.go b/tools/pd-ctl/tests/store/store_test.go index e20ced8e88a..7e04fba10f2 100644 --- a/tools/pd-ctl/tests/store/store_test.go +++ b/tools/pd-ctl/tests/store/store_test.go @@ -66,6 +66,7 @@ func (s *storeTestSuite) TearDownTest() { func (s *storeTestSuite) TestStoreLimitV2() { s.env.RunTest(s.checkStoreLimitV2) } + func (s *storeTestSuite) checkStoreLimitV2(cluster *pdTests.TestCluster) { re := s.Require() pdAddr := cluster.GetConfig().GetClientURL()