From ed1fa0a1801f0677486cfc11cf89c855bb8fd365 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Thu, 28 Nov 2024 16:23:33 +0800 Subject: [PATCH 01/13] reset prepare checker once the cache is reset Signed-off-by: Ryan Leung --- pkg/btree/btree_generic_test.go | 43 ----------- pkg/mcs/scheduling/server/apis/v1/api.go | 1 + pkg/mcs/scheduling/server/cluster.go | 5 ++ pkg/mcs/scheduling/server/server.go | 3 +- pkg/mock/mockcluster/mockcluster.go | 3 + pkg/schedule/checker/checker_controller.go | 11 ++- pkg/schedule/coordinator.go | 36 +++++++-- pkg/schedule/core/cluster_informer.go | 2 +- pkg/schedule/prepare_checker.go | 13 +++- .../schedulers/scheduler_controller.go | 7 +- .../unsafe_recovery_controller.go | 2 + server/api/admin.go | 1 + server/cluster/cluster_test.go | 4 + server/cluster/scheduling_controller.go | 7 ++ .../mcs/scheduling/server_test.go | 1 - tests/server/api/api_test.go | 77 ++++++++++++++++++- tests/server/api/region_test.go | 4 + tests/server/cluster/cluster_test.go | 5 +- .../pd-ctl/pdctl/command/operator_command.go | 4 +- tools/pd-ctl/tests/store/store_test.go | 1 + 20 files changed, 167 insertions(+), 63 deletions(-) diff --git a/pkg/btree/btree_generic_test.go b/pkg/btree/btree_generic_test.go index 114a7238518..e3ee90e16d2 100644 --- a/pkg/btree/btree_generic_test.go +++ b/pkg/btree/btree_generic_test.go @@ -30,7 +30,6 @@ package btree import ( "flag" - "fmt" "math/rand" "reflect" "sort" @@ -222,48 +221,6 @@ func TestBTreeG(t *testing.T) { } } -func ExampleBTreeG() { - tr := NewG[Int](*btreeDegree) - for i := Int(0); i < 10; i++ { - tr.ReplaceOrInsert(i) - } - fmt.Println("len: ", tr.Len()) - v, ok := tr.Get(3) - fmt.Println("get3: ", v, ok) - v, ok = tr.Get(100) - fmt.Println("get100: ", v, ok) - v, ok = tr.Delete(4) - fmt.Println("del4: ", v, ok) - v, ok = tr.Delete(100) - fmt.Println("del100: ", v, ok) - v, ok = tr.ReplaceOrInsert(5) - fmt.Println("replace5: ", v, ok) - v, ok = tr.ReplaceOrInsert(100) - fmt.Println("replace100:", v, ok) - v, ok = tr.Min() - fmt.Println("min: ", v, ok) - v, ok = tr.DeleteMin() - fmt.Println("delmin: ", v, ok) - v, ok = tr.Max() - fmt.Println("max: ", v, ok) - v, ok = tr.DeleteMax() - fmt.Println("delmax: ", v, ok) - fmt.Println("len: ", tr.Len()) - // Output: - // len: 10 - // get3: 3 true - // get100: 0 false - // del4: 4 true - // del100: 0 false - // replace5: 5 true - // replace100: 0 false - // min: 0 true - // delmin: 0 true - // max: 100 true - // delmax: 100 true - // len: 8 -} - func TestDeleteMinG(t *testing.T) { tr := NewG[Int](3) for _, v := range perm(100) { diff --git a/pkg/mcs/scheduling/server/apis/v1/api.go b/pkg/mcs/scheduling/server/apis/v1/api.go index 2efc30c1a18..acffc97b053 100644 --- a/pkg/mcs/scheduling/server/apis/v1/api.go +++ b/pkg/mcs/scheduling/server/apis/v1/api.go @@ -302,6 +302,7 @@ func deleteAllRegionCache(c *gin.Context) { return } cluster.ResetRegionCache() + cluster.ResetPrepared() 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 fa20ef6aba8..12f2c07591d 100644 --- a/pkg/mcs/scheduling/server/cluster.go +++ b/pkg/mcs/scheduling/server/cluster.go @@ -821,6 +821,11 @@ func (c *Cluster) SetPrepared() { c.coordinator.GetPrepareChecker().SetPrepared() } +// ResetPrepared reset the prepare checker. +func (c *Cluster) ResetPrepared() { + c.coordinator.GetPrepareChecker().ResetPrepared() +} + // 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 5c71b5e4f94..802f9155f8b 100644 --- a/pkg/mcs/scheduling/server/server.go +++ b/pkg/mcs/scheduling/server/server.go @@ -16,6 +16,7 @@ package server import ( "context" + "fmt" "net/http" "os" "os/signal" @@ -222,7 +223,7 @@ func (s *Server) updatePDMemberLoop() { } if s.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", fmt.Sprintf("%x", curLeader)), zap.String("new-leader-id", fmt.Sprintf("%x", ep.ID)), 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 6b1105e6a32..ae8b79c647a 100644 --- a/pkg/mock/mockcluster/mockcluster.go +++ b/pkg/mock/mockcluster/mockcluster.go @@ -922,3 +922,6 @@ func (mc *Cluster) ObserveRegionsStats() { storeIDs, writeBytesRates, writeKeysRates := mc.GetStoresWriteRate() mc.HotStat.ObserveRegionsStats(storeIDs, writeBytesRates, writeKeysRates) } + +// ResetPrepared mocks method. +func (*Cluster) ResetPrepared() {} diff --git a/pkg/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index 805da161764..0676289f4ab 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -35,7 +35,6 @@ 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/placement" "github.com/tikv/pd/pkg/utils/keyutil" "github.com/tikv/pd/pkg/utils/logutil" ) @@ -93,12 +92,16 @@ type Controller struct { // patrolRegionScanLimit is the limit of regions to scan. // It is calculated by the number of regions. patrolRegionScanLimit int + prepareChecker *sche.PrepareChecker metrics *checkerControllerMetrics } // NewController create a new Controller. -func NewController(ctx context.Context, cluster sche.CheckerCluster, conf config.CheckerConfigProvider, ruleManager *placement.RuleManager, labeler *labeler.RegionLabeler, opController *operator.Controller) *Controller { +func NewController(ctx context.Context, cluster sche.CheckerCluster, opController *operator.Controller, prepareChecker *sche.PrepareChecker) *Controller { pendingProcessedRegions := cache.NewIDTTL(ctx, time.Minute, 3*time.Minute) + conf := cluster.GetCheckerConfig() + ruleManager := cluster.GetRuleManager() + labeler := cluster.GetRegionLabeler() c := &Controller{ ctx: ctx, cluster: cluster, @@ -116,6 +119,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.duration.Store(time.Duration(0)) @@ -140,6 +144,9 @@ func (c *Controller) PatrolRegions() { case <-ticker.C: c.updateTickerIfNeeded(ticker) c.updatePatrolWorkersIfNeeded() + if !c.prepareChecker.IsPrepared() { + continue + } if c.cluster.IsSchedulingHalted() { for len(c.patrolRegionContext.regionChan) > 0 { <-c.patrolRegionContext.regionChan diff --git a/pkg/schedule/coordinator.go b/pkg/schedule/coordinator.go index 571b56bc4f5..f8bd1e4ef7f 100644 --- a/pkg/schedule/coordinator.go +++ b/pkg/schedule/coordinator.go @@ -46,8 +46,8 @@ import ( const ( runSchedulerCheckInterval = 3 * time.Second - // CollectTimeout is the timeout for collecting regions. - CollectTimeout = 5 * time.Minute + // collectTimeout is the timeout for collecting regions. + collectTimeout = 5 * time.Minute maxLoadConfigRetries = 10 // pushOperatorTickInterval is the interval try to push the operator. pushOperatorTickInterval = 500 * time.Millisecond @@ -83,15 +83,16 @@ 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 := newPrepareChecker(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(), cluster.GetRuleManager(), cluster.GetRegionLabeler(), opController) + schedulers := schedulers.NewController(ctx, cluster, cluster.GetStorage(), opController, prepareChecker) + checkers := checker.NewController(ctx, cluster, opController, prepareChecker) return &Coordinator{ ctx: ctx, 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), @@ -206,6 +207,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() @@ -239,7 +260,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 @@ -551,7 +573,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. diff --git a/pkg/schedule/core/cluster_informer.go b/pkg/schedule/core/cluster_informer.go index dec890ede80..856500e10dd 100644 --- a/pkg/schedule/core/cluster_informer.go +++ b/pkg/schedule/core/cluster_informer.go @@ -41,7 +41,6 @@ type SchedulerCluster interface { buckets.BucketStatInformer GetSchedulerConfig() sc.SchedulerConfigProvider - GetRegionLabeler() *labeler.RegionLabeler GetStoreConfig() sc.StoreConfigProvider } @@ -63,6 +62,7 @@ type SharedCluster interface { GetSharedConfig() sc.SharedConfigProvider GetRuleManager() *placement.RuleManager GetKeyRangeManager() *keyrange.Manager + GetRegionLabeler() *labeler.RegionLabeler AllocID(uint32) (uint64, uint32, error) IsSchedulingHalted() bool GetPrepareRegionCount() (int, error) diff --git a/pkg/schedule/prepare_checker.go b/pkg/schedule/prepare_checker.go index 2ec4c5aaf35..bfda5209028 100644 --- a/pkg/schedule/prepare_checker.go +++ b/pkg/schedule/prepare_checker.go @@ -25,6 +25,7 @@ import ( "github.com/tikv/pd/pkg/utils/syncutil" ) +// prepareChecker is used to check if the coordinator has finished cluster information preparation. type prepareChecker struct { syncutil.RWMutex start time.Time @@ -40,13 +41,13 @@ func newPrepareChecker(totalRegionCountFn func() (int, error)) *prepareChecker { } // 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 { +func (checker *prepareChecker) 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 } @@ -93,3 +94,11 @@ func (checker *prepareChecker) SetPrepared() { defer checker.Unlock() checker.prepared = true } + +// ResetPrepared is for test purpose +func (checker *prepareChecker) ResetPrepared() { + checker.Lock() + defer checker.Unlock() + checker.prepared = false + checker.start = time.Now() +} diff --git a/pkg/schedule/schedulers/scheduler_controller.go b/pkg/schedule/schedulers/scheduler_controller.go index 25c20f2846c..852dad2a84c 100644 --- a/pkg/schedule/schedulers/scheduler_controller.go +++ b/pkg/schedule/schedulers/scheduler_controller.go @@ -57,10 +57,11 @@ type Controller struct { // which will only be initialized and used in the microservice env now. schedulerHandlers map[string]http.Handler opController *operator.Controller + prepareChecker *sche.PrepareChecker } // 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, prepareChecker *sche.PrepareChecker) *Controller { return &Controller{ ctx: ctx, cluster: cluster, @@ -68,6 +69,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, } } @@ -369,6 +371,9 @@ func (c *Controller) runScheduler(s *ScheduleController) { for { select { case <-ticker.C: + if !c.prepareChecker.IsPrepared() { + continue + } diagnosable := s.IsDiagnosticAllowed() if !s.AllowSchedule(diagnosable) { continue diff --git a/pkg/unsaferecovery/unsafe_recovery_controller.go b/pkg/unsaferecovery/unsafe_recovery_controller.go index 40ab186f7e4..df271117448 100644 --- a/pkg/unsaferecovery/unsafe_recovery_controller.go +++ b/pkg/unsaferecovery/unsafe_recovery_controller.go @@ -109,6 +109,7 @@ const ( type cluster interface { core.StoreSetInformer + ResetPrepared() ResetRegionCache() AllocID(uint32) (uint64, uint32, error) BuryStore(storeID uint64, forceBury bool) error @@ -554,6 +555,7 @@ func (u *Controller) changeStage(stage stage) { if u.step > 1 { // == 1 means no operation has done, no need to invalid cache u.cluster.ResetRegionCache() + u.cluster.ResetPrepared() } output.Info = "Unsafe recovery Finished" output.Details = u.getAffectedTableDigest() diff --git a/server/api/admin.go b/server/api/admin.go index 5a5423c1ef7..bc3491d8a98 100644 --- a/server/api/admin.go +++ b/server/api/admin.go @@ -134,6 +134,7 @@ func (h *adminHandler) DeleteAllRegionCache(w http.ResponseWriter, r *http.Reque var err error rc := getCluster(r) rc.ResetRegionCache() + rc.ResetPrepared() msg := "All regions are removed from server cache." if rc.IsServiceIndependent(constant.SchedulingServiceName) { err = h.deleteRegionCacheInSchedulingServer() diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go index 3f327aee18f..8f9df14279a 100644 --- a/server/cluster/cluster_test.go +++ b/server/cluster/cluster_test.go @@ -2879,6 +2879,7 @@ func TestCheckCache(t *testing.T) { cfg.ReplicaScheduleLimit = 0 }, nil, nil, re) defer cleanup() + co.GetPrepareChecker().SetPrepared() oc := co.GetOperatorController() checker := co.GetCheckerController() @@ -2934,6 +2935,7 @@ func TestPatrolRegionConcurrency(t *testing.T) { cfg.MergeScheduleLimit = uint64(mergeScheduleLimit) }, nil, nil, re) defer cleanup() + co.GetPrepareChecker().SetPrepared() oc := co.GetOperatorController() checker := co.GetCheckerController() @@ -2996,6 +2998,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..64875e0879d 100644 --- a/server/cluster/scheduling_controller.go +++ b/server/cluster/scheduling_controller.go @@ -484,6 +484,13 @@ 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() +} + // IsSchedulingControllerRunning returns whether the scheduling controller is running. Only for test purpose. func (sc *schedulingController) IsSchedulingControllerRunning() bool { sc.mu.RLock() diff --git a/tests/integrations/mcs/scheduling/server_test.go b/tests/integrations/mcs/scheduling/server_test.go index 585df59a43e..ba06c9b19b3 100644 --- a/tests/integrations/mcs/scheduling/server_test.go +++ b/tests/integrations/mcs/scheduling/server_test.go @@ -634,7 +634,6 @@ func (suite *serverTestSuite) TestStoreLimit() { re.NoError(err) defer tc.Destroy() tc.WaitForPrimaryServing(re) - oc := tc.GetPrimaryServer().GetCluster().GetCoordinator().GetOperatorController() leaderServer := suite.pdLeader.GetServer() conf := leaderServer.GetReplicationConfig().Clone() diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go index 9be239ffacd..48e8434e68e 100644 --- a/tests/server/api/api_test.go +++ b/tests/server/api/api_test.go @@ -37,6 +37,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" @@ -171,7 +172,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")) @@ -1285,3 +1285,78 @@ 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 + }) + rc.HandleRegionHeartbeat(regionInfo) + 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") + rc.GetOperatorController().RemoveOperators() + re.Equal(0, int(rc.GetOperatorController().OperatorCount(operator.OpSplit))) + + // Simulate continuous heartbeat and verify scheduling recovery + rc.HandleRegionHeartbeat(regionInfo) + + testutil.Eventually(re, func() bool { + count := rc.GetOperatorController().OperatorCount(operator.OpSplit) + return count > 0 + }) +} diff --git a/tests/server/api/region_test.go b/tests/server/api/region_test.go index 52778367601..dd3470df65f 100644 --- a/tests/server/api/region_test.go +++ b/tests/server/api/region_test.go @@ -53,10 +53,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 ed83eb54d5a..2f299f361ab 100644 --- a/tests/server/cluster/cluster_test.go +++ b/tests/server/cluster/cluster_test.go @@ -1927,9 +1927,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() @@ -2061,8 +2064,6 @@ func TestPatrolRegionConfigChange(t *testing.T) { endKey := []byte(strconv.Itoa(i * 2)) tests.MustPutRegion(re, tc, uint64(i), uint64(i%3+1), startKey, endKey) } - fname := testutil.InitTempFileLogger("debug") - defer os.RemoveAll(fname) checkLog(re, fname, "coordinator starts patrol regions") // test change patrol region interval 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 a867ff93af5..655c07542d2 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() From 05feed83aefb3a17ed17c02b018cbeaa9eee573b Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Wed, 15 Oct 2025 12:08:01 +0800 Subject: [PATCH 02/13] resolve conflicts Signed-off-by: Ryan Leung --- pkg/mcs/scheduling/server/server.go | 3 +-- pkg/schedule/checker/checker_controller.go | 5 ++-- pkg/schedule/coordinator.go | 11 ++++---- .../{ => preparecheck}/prepare_checker.go | 26 +++++++++++-------- .../schedulers/scheduler_controller.go | 5 ++-- tests/server/api/api_test.go | 6 +++-- tests/server/cluster/cluster_test.go | 2 ++ 7 files changed, 33 insertions(+), 25 deletions(-) rename pkg/schedule/{ => preparecheck}/prepare_checker.go (78%) diff --git a/pkg/mcs/scheduling/server/server.go b/pkg/mcs/scheduling/server/server.go index 802f9155f8b..38d63dc6933 100644 --- a/pkg/mcs/scheduling/server/server.go +++ b/pkg/mcs/scheduling/server/server.go @@ -16,7 +16,6 @@ package server import ( "context" - "fmt" "net/http" "os" "os/signal" @@ -223,7 +222,7 @@ func (s *Server) updatePDMemberLoop() { } if s.cluster.SwitchPDLeader(pdpb.NewPDClient(cc)) { if status.Leader != curLeader { - log.Info("switch PD leader", zap.String("current-leader", fmt.Sprintf("%x", curLeader)), zap.String("new-leader-id", fmt.Sprintf("%x", ep.ID)), 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/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index 0676289f4ab..460fc3c74e1 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" ) @@ -92,12 +93,12 @@ type Controller struct { // patrolRegionScanLimit is the limit of regions to scan. // It is calculated by the number of regions. patrolRegionScanLimit int - prepareChecker *sche.PrepareChecker + prepareChecker *preparecheck.Checker metrics *checkerControllerMetrics } // NewController create a new Controller. -func NewController(ctx context.Context, cluster sche.CheckerCluster, opController *operator.Controller, prepareChecker *sche.PrepareChecker) *Controller { +func NewController(ctx context.Context, cluster sche.CheckerCluster, opController *operator.Controller, prepareChecker *preparecheck.Checker) *Controller { pendingProcessedRegions := cache.NewIDTTL(ctx, time.Minute, 3*time.Minute) conf := cluster.GetCheckerConfig() ruleManager := cluster.GetRuleManager() diff --git a/pkg/schedule/coordinator.go b/pkg/schedule/coordinator.go index f8bd1e4ef7f..62eb8ad46c4 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,7 +82,7 @@ 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 := newPrepareChecker(cluster.GetPrepareRegionCount) + prepareChecker := preparecheck.NewChecker(cluster.GetPrepareRegionCount) opController := operator.NewController(ctx, cluster.GetBasicCluster(), cluster.GetSharedConfig(), hbStreams) schedulers := schedulers.NewController(ctx, cluster, cluster.GetStorage(), opController, prepareChecker) checkers := checker.NewController(ctx, cluster, opController, prepareChecker) @@ -641,7 +640,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 78% rename from pkg/schedule/prepare_checker.go rename to pkg/schedule/preparecheck/prepare_checker.go index bfda5209028..022da3cc43f 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,23 +25,27 @@ import ( "github.com/tikv/pd/pkg/utils/syncutil" ) -// prepareChecker is used to check if the coordinator has finished cluster information preparation. -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 { @@ -79,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 } @@ -89,14 +93,14 @@ func (checker *prepareChecker) IsPrepared() bool { } // SetPrepared is for test purpose -func (checker *prepareChecker) SetPrepared() { +func (checker *Checker) SetPrepared() { checker.Lock() defer checker.Unlock() checker.prepared = true } // ResetPrepared is for test purpose -func (checker *prepareChecker) ResetPrepared() { +func (checker *Checker) ResetPrepared() { checker.Lock() defer checker.Unlock() checker.prepared = false diff --git a/pkg/schedule/schedulers/scheduler_controller.go b/pkg/schedule/schedulers/scheduler_controller.go index 852dad2a84c..991cb9fa180 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,11 +58,11 @@ type Controller struct { // which will only be initialized and used in the microservice env now. schedulerHandlers map[string]http.Handler opController *operator.Controller - prepareChecker *sche.PrepareChecker + prepareChecker *preparecheck.Checker } // NewController creates a scheduler controller. -func NewController(ctx context.Context, cluster sche.SchedulerCluster, storage endpoint.ConfigStorage, opController *operator.Controller, prepareChecker *sche.PrepareChecker) *Controller { +func NewController(ctx context.Context, cluster sche.SchedulerCluster, storage endpoint.ConfigStorage, opController *operator.Controller, prepareChecker *preparecheck.Checker) *Controller { return &Controller{ ctx: ctx, cluster: cluster, diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go index 48e8434e68e..8507465da65 100644 --- a/tests/server/api/api_test.go +++ b/tests/server/api/api_test.go @@ -1335,7 +1335,8 @@ func TestDeleteAllRegionCacheScheduling(t *testing.T) { testutil.Eventually(re, func() bool { return rc.GetCoordinator() != nil }) - rc.HandleRegionHeartbeat(regionInfo) + err = rc.HandleRegionHeartbeat(regionInfo) + re.NoError(err) testutil.Eventually(re, func() bool { return rc.GetCoordinator().GetPrepareChecker().IsPrepared() }) @@ -1353,7 +1354,8 @@ func TestDeleteAllRegionCacheScheduling(t *testing.T) { re.Equal(0, int(rc.GetOperatorController().OperatorCount(operator.OpSplit))) // Simulate continuous heartbeat and verify scheduling recovery - rc.HandleRegionHeartbeat(regionInfo) + err = rc.HandleRegionHeartbeat(regionInfo) + re.NoError(err) testutil.Eventually(re, func() bool { count := rc.GetOperatorController().OperatorCount(operator.OpSplit) diff --git a/tests/server/cluster/cluster_test.go b/tests/server/cluster/cluster_test.go index 2f299f361ab..c20255c1dc7 100644 --- a/tests/server/cluster/cluster_test.go +++ b/tests/server/cluster/cluster_test.go @@ -2064,6 +2064,8 @@ func TestPatrolRegionConfigChange(t *testing.T) { endKey := []byte(strconv.Itoa(i * 2)) tests.MustPutRegion(re, tc, uint64(i), uint64(i%3+1), startKey, endKey) } + fname := testutil.InitTempFileLogger("debug") + defer os.RemoveAll(fname) checkLog(re, fname, "coordinator starts patrol regions") // test change patrol region interval From 390bbc326d097020a63f4bcaa360f6a222b2c15e Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Wed, 15 Oct 2025 13:39:23 +0800 Subject: [PATCH 03/13] fix Signed-off-by: Ryan Leung --- tests/server/api/diagnostic_test.go | 2 ++ 1 file changed, 2 insertions(+) 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) From 8c5262edd84621e59d997359d058bf9bc627475e Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 16:58:21 +0800 Subject: [PATCH 04/13] tests: mark split scatter cluster prepared Signed-off-by: Ryan Leung --- server/cluster/split_scatter_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/server/cluster/split_scatter_test.go b/server/cluster/split_scatter_test.go index a6e0793631f..6cb6cac0255 100644 --- a/server/cluster/split_scatter_test.go +++ b/server/cluster/split_scatter_test.go @@ -168,6 +168,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() }) From 81dcb1390087ae43bd6f23fc8fcb39557cbef6c5 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 17:32:39 +0800 Subject: [PATCH 05/13] tests: wait for read hot region history Signed-off-by: Ryan Leung --- tests/server/storage/hot_region_storage_test.go | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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) From 7cab94e56b504e0f18ff87e63d1149ecc1113053 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 17:45:51 +0800 Subject: [PATCH 06/13] schedule: reset prepare checker before clearing cache Signed-off-by: Ryan Leung --- pkg/mcs/scheduling/server/apis/v1/api.go | 2 +- pkg/schedule/preparecheck/prepare_checker.go | 3 ++- pkg/unsaferecovery/unsafe_recovery_controller.go | 2 +- server/api/admin.go | 2 +- tests/server/api/api_test.go | 1 + 5 files changed, 6 insertions(+), 4 deletions(-) diff --git a/pkg/mcs/scheduling/server/apis/v1/api.go b/pkg/mcs/scheduling/server/apis/v1/api.go index 49b13f57b3a..1d62314c725 100644 --- a/pkg/mcs/scheduling/server/apis/v1/api.go +++ b/pkg/mcs/scheduling/server/apis/v1/api.go @@ -383,8 +383,8 @@ func deleteAllRegionCache(c *gin.Context) { c.String(http.StatusInternalServerError, errs.ErrNotBootstrapped.GenWithStackByArgs().Error()) return } - cluster.ResetRegionCache() cluster.ResetPrepared() + cluster.ResetRegionCache() c.String(http.StatusOK, "All regions are removed from server cache.") } diff --git a/pkg/schedule/preparecheck/prepare_checker.go b/pkg/schedule/preparecheck/prepare_checker.go index 022da3cc43f..424ab73a2bb 100644 --- a/pkg/schedule/preparecheck/prepare_checker.go +++ b/pkg/schedule/preparecheck/prepare_checker.go @@ -99,7 +99,8 @@ func (checker *Checker) SetPrepared() { checker.prepared = true } -// ResetPrepared is for test purpose +// ResetPrepared resets the checker so scheduling pauses until cluster +// information is collected again after the region cache is reset. func (checker *Checker) ResetPrepared() { checker.Lock() defer checker.Unlock() diff --git a/pkg/unsaferecovery/unsafe_recovery_controller.go b/pkg/unsaferecovery/unsafe_recovery_controller.go index 0959c759155..a3891cabd98 100644 --- a/pkg/unsaferecovery/unsafe_recovery_controller.go +++ b/pkg/unsaferecovery/unsafe_recovery_controller.go @@ -587,8 +587,8 @@ func (u *Controller) changeStage(stage stage) { case Finished: if u.step > 1 { // == 1 means no operation has done, no need to invalid cache - u.cluster.ResetRegionCache() u.cluster.ResetPrepared() + u.cluster.ResetRegionCache() } output.Info = "Unsafe recovery Finished" output.Details = u.getAffectedTableDigest() diff --git a/server/api/admin.go b/server/api/admin.go index 4f8a0b7fada..54cf03efefc 100644 --- a/server/api/admin.go +++ b/server/api/admin.go @@ -154,8 +154,8 @@ 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() rc.ResetPrepared() + rc.ResetRegionCache() msg := "All regions are removed from server cache." if rc.IsServiceIndependent(constant.SchedulingServiceName) { err = h.deleteRegionCacheInSchedulingServer() diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go index 92861a88277..9a8964e059e 100644 --- a/tests/server/api/api_test.go +++ b/tests/server/api/api_test.go @@ -1759,6 +1759,7 @@ func TestDeleteAllRegionCacheScheduling(t *testing.T) { 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))) From c4b3916961b2fe064976eab5d0d1bd124ffd395a Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 18:28:40 +0800 Subject: [PATCH 07/13] schedule: clear operators when resetting prepare checker Signed-off-by: Ryan Leung --- pkg/mcs/scheduling/server/cluster.go | 1 + pkg/schedule/operator/operator_controller.go | 25 +++++++++++++++++++ .../operator/operator_controller_test.go | 22 ++++++++++++++++ pkg/schedule/operator/waiting_operator.go | 21 ++++++++++++++++ server/cluster/scheduling_controller.go | 4 +++ tests/server/api/api_test.go | 2 +- 6 files changed, 74 insertions(+), 1 deletion(-) diff --git a/pkg/mcs/scheduling/server/cluster.go b/pkg/mcs/scheduling/server/cluster.go index 417a6e13e28..2898b0ec720 100644 --- a/pkg/mcs/scheduling/server/cluster.go +++ b/pkg/mcs/scheduling/server/cluster.go @@ -942,6 +942,7 @@ func (c *Cluster) SetPrepared() { // ResetPrepared reset the prepare checker. func (c *Cluster) ResetPrepared() { c.coordinator.GetPrepareChecker().ResetPrepared() + c.coordinator.GetOperatorController().RemoveAllOperators(operator.AdminStop) } // IsSchedulingHalted returns whether the scheduling is halted. diff --git a/pkg/schedule/operator/operator_controller.go b/pkg/schedule/operator/operator_controller.go index 80cffcc399b..a2cc60f5707 100644 --- a/pkg/schedule/operator/operator_controller.go +++ b/pkg/schedule/operator/operator_controller.go @@ -648,6 +648,31 @@ func (oc *Controller) RemoveOperators(reasons ...CancelReasonType) { } } +// RemoveAllOperators removes all running and waiting operators. +func (oc *Controller) RemoveAllOperators(reasons ...CancelReasonType) { + oc.RemoveOperators(reasons...) + + waiting := oc.wop.Clear() + oc.wopStatus.clear() + var cancelReason CancelReasonType + if len(reasons) > 0 { + cancelReason = reasons[0] + } + for _, op := range waiting { + if op == nil { + continue + } + operatorCounter.WithLabelValues(op.Desc(), "remove").Inc() + if op.Cancel(cancelReason) { + log.Info("waiting operator removed", + zap.Uint64("region-id", op.RegionID()), + zap.Duration("lives", op.ElapsedTime()), + zap.Reflect("operator", op)) + } + oc.buryOperator(op) + } +} + func (oc *Controller) removeOperatorsWithoutBury() []*Operator { var removed []*Operator oc.operators.Range(func(regionID, value any) bool { diff --git a/pkg/schedule/operator/operator_controller_test.go b/pkg/schedule/operator/operator_controller_test.go index c2e7157e2b8..656fe88f7d8 100644 --- a/pkg/schedule/operator/operator_controller_test.go +++ b/pkg/schedule/operator/operator_controller_test.go @@ -953,6 +953,28 @@ func checkRemoveOperatorSuccess(re *require.Assertions, oc *Controller, op *Oper re.Equal(op, oc.GetOperatorStatus(op.RegionID()).Operator) } +func (suite *operatorControllerTestSuite) TestRemoveAllOperatorsClearsWaitingOperators() { + re := suite.Require() + opts := mockconfig.NewTestOptions() + cluster := mockcluster.NewCluster(suite.ctx, opts) + stream := hbstream.NewTestHeartbeatStreams(suite.ctx, cluster, false /* no need to run */) + controller := NewController(suite.ctx, cluster.GetBasicCluster(), cluster.GetSharedConfig(), stream) + + op1 := NewTestOperator(1, &metapb.RegionEpoch{}, OpLeader, TransferLeader{ToStore: 2}) + op2 := NewTestOperator(2, &metapb.RegionEpoch{}, OpLeader, TransferLeader{ToStore: 2}) + controller.wop.PutOperator(op1) + controller.wop.PutOperator(op2) + controller.wopStatus.incCount(op1.Desc()) + controller.wopStatus.incCount(op2.Desc()) + re.Len(controller.GetWaitingOperators(), 2) + + controller.RemoveAllOperators(AdminStop) + re.Empty(controller.GetOperators()) + re.Empty(controller.GetWaitingOperators()) + re.True(op1.IsEnd()) + re.True(op2.IsEnd()) +} + func (suite *operatorControllerTestSuite) TestAddWaitingOperator() { re := suite.Require() opts := mockconfig.NewTestOptions() diff --git a/pkg/schedule/operator/waiting_operator.go b/pkg/schedule/operator/waiting_operator.go index 70013e143c7..ca5e4cedf3c 100644 --- a/pkg/schedule/operator/waiting_operator.go +++ b/pkg/schedule/operator/waiting_operator.go @@ -29,6 +29,7 @@ type WaitingOperator interface { PutMergeOperators(op []*Operator) GetOperator() []*Operator ListOperator() []*Operator + Clear() []*Operator } // bucket is used to maintain the operators created by a specific scheduler. @@ -94,6 +95,20 @@ func (b *randBuckets) ListOperator() []*Operator { return ops } +// Clear removes all operators from the random buckets and returns them. +func (b *randBuckets) Clear() []*Operator { + b.mu.Lock() + defer b.mu.Unlock() + var ops []*Operator + for i := range b.buckets { + bucket := b.buckets[i] + ops = append(ops, bucket.ops...) + bucket.ops = nil + } + b.totalWeight = 0 + return ops +} + // GetOperator gets an operator from the random buckets. func (b *randBuckets) GetOperator() []*Operator { b.mu.Lock() @@ -165,3 +180,9 @@ func (s *waitingOperatorStatus) getCount(kind string) uint64 { defer s.mu.Unlock() return s.ops[kind] } + +func (s *waitingOperatorStatus) clear() { + s.mu.Lock() + defer s.mu.Unlock() + s.ops = make(map[string]uint64) +} diff --git a/server/cluster/scheduling_controller.go b/server/cluster/scheduling_controller.go index 64875e0879d..d1b060cc198 100644 --- a/server/cluster/scheduling_controller.go +++ b/server/cluster/scheduling_controller.go @@ -488,7 +488,11 @@ func (sc *schedulingController) SetPrepared() { func (sc *schedulingController) ResetPrepared() { sc.mu.RLock() defer sc.mu.RUnlock() + if sc.coordinator == nil { + return + } sc.coordinator.GetPrepareChecker().ResetPrepared() + sc.coordinator.GetOperatorController().RemoveAllOperators(operator.AdminStop) } // IsSchedulingControllerRunning returns whether the scheduling controller is running. Only for test purpose. diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go index 9a8964e059e..8b887b4fb4b 100644 --- a/tests/server/api/api_test.go +++ b/tests/server/api/api_test.go @@ -1760,8 +1760,8 @@ func TestDeleteAllRegionCacheScheduling(t *testing.T) { 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))) + re.Empty(rc.GetOperatorController().GetWaitingOperators()) // Simulate continuous heartbeat and verify scheduling recovery err = rc.HandleRegionHeartbeat(regionInfo) From d13a2f8ab958b52d853bec59945f367b7f9dae2d Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 18:45:08 +0800 Subject: [PATCH 08/13] Revert "schedule: clear operators when resetting prepare checker" This reverts commit c4b3916961b2fe064976eab5d0d1bd124ffd395a. Signed-off-by: Ryan Leung --- pkg/mcs/scheduling/server/cluster.go | 1 - pkg/schedule/operator/operator_controller.go | 25 ------------------- .../operator/operator_controller_test.go | 22 ---------------- pkg/schedule/operator/waiting_operator.go | 21 ---------------- server/cluster/scheduling_controller.go | 4 --- tests/server/api/api_test.go | 2 +- 6 files changed, 1 insertion(+), 74 deletions(-) diff --git a/pkg/mcs/scheduling/server/cluster.go b/pkg/mcs/scheduling/server/cluster.go index 2898b0ec720..417a6e13e28 100644 --- a/pkg/mcs/scheduling/server/cluster.go +++ b/pkg/mcs/scheduling/server/cluster.go @@ -942,7 +942,6 @@ func (c *Cluster) SetPrepared() { // ResetPrepared reset the prepare checker. func (c *Cluster) ResetPrepared() { c.coordinator.GetPrepareChecker().ResetPrepared() - c.coordinator.GetOperatorController().RemoveAllOperators(operator.AdminStop) } // IsSchedulingHalted returns whether the scheduling is halted. diff --git a/pkg/schedule/operator/operator_controller.go b/pkg/schedule/operator/operator_controller.go index a2cc60f5707..80cffcc399b 100644 --- a/pkg/schedule/operator/operator_controller.go +++ b/pkg/schedule/operator/operator_controller.go @@ -648,31 +648,6 @@ func (oc *Controller) RemoveOperators(reasons ...CancelReasonType) { } } -// RemoveAllOperators removes all running and waiting operators. -func (oc *Controller) RemoveAllOperators(reasons ...CancelReasonType) { - oc.RemoveOperators(reasons...) - - waiting := oc.wop.Clear() - oc.wopStatus.clear() - var cancelReason CancelReasonType - if len(reasons) > 0 { - cancelReason = reasons[0] - } - for _, op := range waiting { - if op == nil { - continue - } - operatorCounter.WithLabelValues(op.Desc(), "remove").Inc() - if op.Cancel(cancelReason) { - log.Info("waiting operator removed", - zap.Uint64("region-id", op.RegionID()), - zap.Duration("lives", op.ElapsedTime()), - zap.Reflect("operator", op)) - } - oc.buryOperator(op) - } -} - func (oc *Controller) removeOperatorsWithoutBury() []*Operator { var removed []*Operator oc.operators.Range(func(regionID, value any) bool { diff --git a/pkg/schedule/operator/operator_controller_test.go b/pkg/schedule/operator/operator_controller_test.go index 656fe88f7d8..c2e7157e2b8 100644 --- a/pkg/schedule/operator/operator_controller_test.go +++ b/pkg/schedule/operator/operator_controller_test.go @@ -953,28 +953,6 @@ func checkRemoveOperatorSuccess(re *require.Assertions, oc *Controller, op *Oper re.Equal(op, oc.GetOperatorStatus(op.RegionID()).Operator) } -func (suite *operatorControllerTestSuite) TestRemoveAllOperatorsClearsWaitingOperators() { - re := suite.Require() - opts := mockconfig.NewTestOptions() - cluster := mockcluster.NewCluster(suite.ctx, opts) - stream := hbstream.NewTestHeartbeatStreams(suite.ctx, cluster, false /* no need to run */) - controller := NewController(suite.ctx, cluster.GetBasicCluster(), cluster.GetSharedConfig(), stream) - - op1 := NewTestOperator(1, &metapb.RegionEpoch{}, OpLeader, TransferLeader{ToStore: 2}) - op2 := NewTestOperator(2, &metapb.RegionEpoch{}, OpLeader, TransferLeader{ToStore: 2}) - controller.wop.PutOperator(op1) - controller.wop.PutOperator(op2) - controller.wopStatus.incCount(op1.Desc()) - controller.wopStatus.incCount(op2.Desc()) - re.Len(controller.GetWaitingOperators(), 2) - - controller.RemoveAllOperators(AdminStop) - re.Empty(controller.GetOperators()) - re.Empty(controller.GetWaitingOperators()) - re.True(op1.IsEnd()) - re.True(op2.IsEnd()) -} - func (suite *operatorControllerTestSuite) TestAddWaitingOperator() { re := suite.Require() opts := mockconfig.NewTestOptions() diff --git a/pkg/schedule/operator/waiting_operator.go b/pkg/schedule/operator/waiting_operator.go index ca5e4cedf3c..70013e143c7 100644 --- a/pkg/schedule/operator/waiting_operator.go +++ b/pkg/schedule/operator/waiting_operator.go @@ -29,7 +29,6 @@ type WaitingOperator interface { PutMergeOperators(op []*Operator) GetOperator() []*Operator ListOperator() []*Operator - Clear() []*Operator } // bucket is used to maintain the operators created by a specific scheduler. @@ -95,20 +94,6 @@ func (b *randBuckets) ListOperator() []*Operator { return ops } -// Clear removes all operators from the random buckets and returns them. -func (b *randBuckets) Clear() []*Operator { - b.mu.Lock() - defer b.mu.Unlock() - var ops []*Operator - for i := range b.buckets { - bucket := b.buckets[i] - ops = append(ops, bucket.ops...) - bucket.ops = nil - } - b.totalWeight = 0 - return ops -} - // GetOperator gets an operator from the random buckets. func (b *randBuckets) GetOperator() []*Operator { b.mu.Lock() @@ -180,9 +165,3 @@ func (s *waitingOperatorStatus) getCount(kind string) uint64 { defer s.mu.Unlock() return s.ops[kind] } - -func (s *waitingOperatorStatus) clear() { - s.mu.Lock() - defer s.mu.Unlock() - s.ops = make(map[string]uint64) -} diff --git a/server/cluster/scheduling_controller.go b/server/cluster/scheduling_controller.go index d1b060cc198..64875e0879d 100644 --- a/server/cluster/scheduling_controller.go +++ b/server/cluster/scheduling_controller.go @@ -488,11 +488,7 @@ func (sc *schedulingController) SetPrepared() { func (sc *schedulingController) ResetPrepared() { sc.mu.RLock() defer sc.mu.RUnlock() - if sc.coordinator == nil { - return - } sc.coordinator.GetPrepareChecker().ResetPrepared() - sc.coordinator.GetOperatorController().RemoveAllOperators(operator.AdminStop) } // IsSchedulingControllerRunning returns whether the scheduling controller is running. Only for test purpose. diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go index 8b887b4fb4b..9a8964e059e 100644 --- a/tests/server/api/api_test.go +++ b/tests/server/api/api_test.go @@ -1760,8 +1760,8 @@ func TestDeleteAllRegionCacheScheduling(t *testing.T) { 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))) - re.Empty(rc.GetOperatorController().GetWaitingOperators()) // Simulate continuous heartbeat and verify scheduling recovery err = rc.HandleRegionHeartbeat(regionInfo) From a778ebdd8b4a960a04349c801e1c54c75f2618d1 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Tue, 7 Jul 2026 18:05:53 +0800 Subject: [PATCH 09/13] schedule: gate queued operators with prepare checker Signed-off-by: Ryan Leung --- pkg/schedule/checker/checker_controller.go | 46 ++++++++------- .../checker/checker_controller_test.go | 57 +++++++++++++++++++ pkg/schedule/preparecheck/prepare_checker.go | 14 +++++ .../preparecheck/prepare_checker_test.go | 47 +++++++++++++++ .../schedulers/scheduler_controller.go | 20 ++++--- 5 files changed, 158 insertions(+), 26 deletions(-) create mode 100644 pkg/schedule/checker/checker_controller_test.go create mode 100644 pkg/schedule/preparecheck/prepare_checker_test.go diff --git a/pkg/schedule/checker/checker_controller.go b/pkg/schedule/checker/checker_controller.go index 5ecaa8a643b..6685242ad86 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -284,27 +284,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. @@ -406,6 +408,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/preparecheck/prepare_checker.go b/pkg/schedule/preparecheck/prepare_checker.go index 424ab73a2bb..f539aa984d6 100644 --- a/pkg/schedule/preparecheck/prepare_checker.go +++ b/pkg/schedule/preparecheck/prepare_checker.go @@ -92,6 +92,20 @@ func (checker *Checker) 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 *Checker) SetPrepared() { checker.Lock() diff --git a/pkg/schedule/preparecheck/prepare_checker_test.go b/pkg/schedule/preparecheck/prepare_checker_test.go new file mode 100644 index 00000000000..093b6f77098 --- /dev/null +++ b/pkg/schedule/preparecheck/prepare_checker_test.go @@ -0,0 +1,47 @@ +// 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 ( + "testing" + + "github.com/stretchr/testify/require" +) + +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) +} diff --git a/pkg/schedule/schedulers/scheduler_controller.go b/pkg/schedule/schedulers/scheduler_controller.go index f7b63967760..9b299726c27 100644 --- a/pkg/schedule/schedulers/scheduler_controller.go +++ b/pkg/schedule/schedulers/scheduler_controller.go @@ -390,17 +390,23 @@ func (c *Controller) runScheduler(s *ScheduleController) { for { select { case <-ticker.C: - if !c.prepareChecker.IsPrepared() { + 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 } - diagnosable := s.IsDiagnosticAllowed() - if !s.AllowSchedule(diagnosable) { + if !shouldResetTicker { 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())) - } // Note: we reset the ticker here to support updating configuration dynamically. ticker.Reset(s.GetInterval()) case <-s.Ctx().Done(): From a13e53b1afc009aa329e3b6396d6a46eba38c782 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Wed, 8 Jul 2026 10:46:11 +0800 Subject: [PATCH 10/13] tests: cover prepare checker test with goleak Signed-off-by: Ryan Leung --- pkg/schedule/preparecheck/prepare_checker_test.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pkg/schedule/preparecheck/prepare_checker_test.go b/pkg/schedule/preparecheck/prepare_checker_test.go index 093b6f77098..58da84c6c0d 100644 --- a/pkg/schedule/preparecheck/prepare_checker_test.go +++ b/pkg/schedule/preparecheck/prepare_checker_test.go @@ -18,8 +18,15 @@ import ( "testing" "github.com/stretchr/testify/require" + "go.uber.org/goleak" + + "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) { From f173c58d42cc8b2a7b96e9274724919529d401ab Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Wed, 8 Jul 2026 11:13:28 +0800 Subject: [PATCH 11/13] tests: wait for affinity groups to sync Signed-off-by: Ryan Leung --- tests/server/apiv2/handlers/affinity_test.go | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/server/apiv2/handlers/affinity_test.go b/tests/server/apiv2/handlers/affinity_test.go index 99f2c4c1feb..78a7de8333a 100644 --- a/tests/server/apiv2/handlers/affinity_test.go +++ b/tests/server/apiv2/handlers/affinity_test.go @@ -856,8 +856,14 @@ 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) + testutil.Eventually(re, func() bool { + result = handlers.AffinityGroupsResponse{} + err := testutil.ReadGetJSON(re, tests.TestDialClient, getAffinityGroupURL(serverAddr)+"?ids=group-1&ids=group-3&ids=missing", &result) + re.NoError(err) + _, ok1 := result.AffinityGroups["group-1"] + _, ok3 := result.AffinityGroups["group-3"] + return len(result.AffinityGroups) == 2 && ok1 && ok3 + }) re.Len(result.AffinityGroups, 2) re.Contains(result.AffinityGroups, "group-1") re.Contains(result.AffinityGroups, "group-3") From 6a3ca64583b70d3c593b05846817f6aa93bd033a Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Wed, 8 Jul 2026 13:16:59 +0800 Subject: [PATCH 12/13] tests: wait for affinity skip-exist list sync Signed-off-by: Ryan Leung --- tests/server/apiv2/handlers/affinity_test.go | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/server/apiv2/handlers/affinity_test.go b/tests/server/apiv2/handlers/affinity_test.go index 78a7de8333a..4d1cb84bb69 100644 --- a/tests/server/apiv2/handlers/affinity_test.go +++ b/tests/server/apiv2/handlers/affinity_test.go @@ -660,7 +660,13 @@ func (suite *affinityHandlerTestSuite) TestAffinityGroupCreateSkipExistCheck() { re.Contains(result.AffinityGroups, "existing") re.Contains(result.AffinityGroups, "new") - listResp := mustGetAllAffinityGroups(re, serverAddr) + var listResp *handlers.AffinityGroupsResponse + testutil.Eventually(re, func() bool { + listResp = mustGetAllAffinityGroups(re, serverAddr) + _, okExisting := listResp.AffinityGroups["existing"] + _, okNew := listResp.AffinityGroups["new"] + return okExisting && okNew + }) re.Contains(listResp.AffinityGroups, "existing") re.Contains(listResp.AffinityGroups, "new") }) From 70b1b5590fa36717878d2bda5820c5bf512e9af3 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 16:57:17 +0800 Subject: [PATCH 13/13] fix: fence scheduling during region cache reset Signed-off-by: Ryan Leung --- pkg/btree/btree_generic_test.go | 43 +++++++ pkg/mcs/scheduling/server/apis/v1/api.go | 3 +- pkg/mcs/scheduling/server/cluster.go | 6 + pkg/mock/mockcluster/mockcluster.go | 11 ++ pkg/schedule/checker/checker_controller.go | 6 +- pkg/schedule/checker/split_scatter_test.go | 31 ++++- pkg/schedule/preparecheck/prepare_checker.go | 10 ++ .../preparecheck/prepare_checker_test.go | 53 ++++++++ .../unsafe_recovery_controller.go | 104 +++++++++++++--- .../unsafe_recovery_controller_test.go | 114 ++++++++++++++++++ server/api/admin.go | 35 +----- server/cluster/cluster.go | 23 +++- server/cluster/scheduling_controller.go | 6 + server/cluster/split_scatter_test.go | 18 +++ server/server.go | 26 ++++ server/server_test.go | 30 +++++ 16 files changed, 461 insertions(+), 58 deletions(-) diff --git a/pkg/btree/btree_generic_test.go b/pkg/btree/btree_generic_test.go index ce4583aa2c8..4d36ef2a7cb 100644 --- a/pkg/btree/btree_generic_test.go +++ b/pkg/btree/btree_generic_test.go @@ -30,6 +30,7 @@ package btree import ( "flag" + "fmt" "math/rand/v2" "reflect" "sort" @@ -221,6 +222,48 @@ func TestBTreeG(t *testing.T) { } } +func ExampleBTreeG() { + tr := NewG[Int](*btreeDegree) + for i := Int(0); i < 10; i++ { + tr.ReplaceOrInsert(i) + } + fmt.Println("len: ", tr.Len()) + v, ok := tr.Get(3) + fmt.Println("get3: ", v, ok) + v, ok = tr.Get(100) + fmt.Println("get100: ", v, ok) + v, ok = tr.Delete(4) + fmt.Println("del4: ", v, ok) + v, ok = tr.Delete(100) + fmt.Println("del100: ", v, ok) + v, ok = tr.ReplaceOrInsert(5) + fmt.Println("replace5: ", v, ok) + v, ok = tr.ReplaceOrInsert(100) + fmt.Println("replace100:", v, ok) + v, ok = tr.Min() + fmt.Println("min: ", v, ok) + v, ok = tr.DeleteMin() + fmt.Println("delmin: ", v, ok) + v, ok = tr.Max() + fmt.Println("max: ", v, ok) + v, ok = tr.DeleteMax() + fmt.Println("delmax: ", v, ok) + fmt.Println("len: ", tr.Len()) + // Output: + // len: 10 + // get3: 3 true + // get100: 0 false + // del4: 4 true + // del100: 0 false + // replace5: 5 true + // replace100: 0 false + // min: 0 true + // delmin: 0 true + // max: 100 true + // delmax: 100 true + // len: 8 +} + func TestDeleteMinG(t *testing.T) { tr := NewG[Int](3) for _, v := range perm(100) { diff --git a/pkg/mcs/scheduling/server/apis/v1/api.go b/pkg/mcs/scheduling/server/apis/v1/api.go index 1d62314c725..b860d9a1a50 100644 --- a/pkg/mcs/scheduling/server/apis/v1/api.go +++ b/pkg/mcs/scheduling/server/apis/v1/api.go @@ -383,8 +383,7 @@ func deleteAllRegionCache(c *gin.Context) { c.String(http.StatusInternalServerError, errs.ErrNotBootstrapped.GenWithStackByArgs().Error()) return } - cluster.ResetPrepared() - 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 417a6e13e28..139a48457be 100644 --- a/pkg/mcs/scheduling/server/cluster.go +++ b/pkg/mcs/scheduling/server/cluster.go @@ -944,6 +944,12 @@ 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/mock/mockcluster/mockcluster.go b/pkg/mock/mockcluster/mockcluster.go index fb1c8038ea3..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() @@ -940,3 +945,9 @@ func (mc *Cluster) ObserveRegionsStats() { // 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 6685242ad86..4001a50ba54 100644 --- a/pkg/schedule/checker/checker_controller.go +++ b/pkg/schedule/checker/checker_controller.go @@ -186,7 +186,7 @@ func (c *Controller) PatrolRegions() { }) measure(c.metrics.patrolPhaseHistograms[phaseDispatchSplitScatter], func() { - c.splitScatter.dispatchSplitScatterRegions() + c.dispatchSplitScatterRegions() }) measure(c.metrics.patrolPhaseHistograms[phaseScanRegions], func() { @@ -223,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() diff --git a/pkg/schedule/checker/split_scatter_test.go b/pkg/schedule/checker/split_scatter_test.go index 19bc3f31e99..c573e3b4c80 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/preparecheck/prepare_checker.go b/pkg/schedule/preparecheck/prepare_checker.go index f539aa984d6..a2e91dea87e 100644 --- a/pkg/schedule/preparecheck/prepare_checker.go +++ b/pkg/schedule/preparecheck/prepare_checker.go @@ -116,8 +116,18 @@ func (checker *Checker) SetPrepared() { // 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 index 58da84c6c0d..9a6893ffbac 100644 --- a/pkg/schedule/preparecheck/prepare_checker_test.go +++ b/pkg/schedule/preparecheck/prepare_checker_test.go @@ -15,11 +15,14 @@ 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" ) @@ -52,3 +55,53 @@ func TestRunIfPrepared(t *testing.T) { })) 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/unsaferecovery/unsafe_recovery_controller.go b/pkg/unsaferecovery/unsafe_recovery_controller.go index be9089f9cb9..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,8 +119,8 @@ const ( type cluster interface { core.StoreSetInformer - ResetPrepared() - ResetRegionCache() + Context() context.Context + ResetPreparedAndResetRegionCache(context.Context) error AllocID(uint32) (uint64, uint32, error) BuryStore(storeID uint64, forceBury bool) error GetSchedulerConfig() sc.SchedulerConfigProvider @@ -158,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. @@ -195,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 @@ -293,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 } @@ -322,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 ( @@ -351,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 } @@ -389,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 } @@ -418,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) { @@ -495,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) } @@ -613,13 +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.ResetPrepared() - 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 54cf03efefc..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,14 +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.ResetPrepared() - 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) @@ -281,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 6ca3b4f57d9..42c8f89676a 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. @@ -2670,6 +2671,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/scheduling_controller.go b/server/cluster/scheduling_controller.go index 64875e0879d..a69d553f286 100644 --- a/server/cluster/scheduling_controller.go +++ b/server/cluster/scheduling_controller.go @@ -491,6 +491,12 @@ func (sc *schedulingController) ResetPrepared() { 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 6cb6cac0255..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) diff --git a/server/server.go b/server/server.go index fe0a2a37d1c..231f98efa02 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, @@ -2315,6 +2318,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()