scheduler: graceful shutdown implement (#9720)

close #9719

Add an is_stopping status to the StoreHeartbeat message. When TiKV receives a SIGTERM, it sets this flag. This change adds a new evict-stopping-store-scheduler to PD, which is analogous to the evict-slow-store-scheduler. It proactively transfers leaders away from nodes by inspecting the is_stopping status from store heartbeats.

Signed-off-by: hujiatao0 <hhjjtt110@gmail.com>

Author: hujiatao0

go.mod

@@ -35,7 +35,7 @@ require (
 	github.com/pingcap/errcode v0.3.0
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86
-	github.com/pingcap/kvproto v0.0.0-20250923064352-8eeada0a8a03
+	github.com/pingcap/kvproto v0.0.0-20250923091925-d79d11002599
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/pingcap/metering_sdk v0.0.0-20250918015914-468cd6feb1dc
 	github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21

go.sum

@@ -469,8 +469,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ue
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20250923064352-8eeada0a8a03 h1:G6lEpMBW42aNvstzIjNR3NnH6mx+3nIH42Ic1Sb8h/U=
-github.com/pingcap/kvproto v0.0.0-20250923064352-8eeada0a8a03/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20250923091925-d79d11002599 h1:57fBeBND/j/dp7nVlw+cWEwYlt4u8CAe4ApsmAEb1ow=
+github.com/pingcap/kvproto v0.0.0-20250923091925-d79d11002599/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=

pkg/core/basic_cluster.go

@@ -149,7 +149,9 @@ type StoreSetController interface {
 	ResumeLeaderTransfer(id uint64, d constant.Direction)
 
 	SlowStoreEvicted(id uint64) error
+	StoppingStoreEvicted(id uint64) error
 	SlowStoreRecovered(id uint64)
+	StoppingStoreRecovered(id uint64)
 	SlowTrendEvicted(id uint64) error
 	SlowTrendRecovered(id uint64)
 }

pkg/core/store.go

@@ -206,6 +206,11 @@ func (s *StoreInfo) EvictedAsSlowStore() bool {
 	return s.slowStoreEvicted.Load() > 0
 }
 
+// EvictedAsStoppingStore returns if the store should be evicted as a stopping store.
+func (s *StoreInfo) EvictedAsStoppingStore() bool {
+	return s.rawStats.IsStopping
+}
+
 // IsEvictedAsSlowTrend returns if the store should be evicted as a slow store by trend.
 func (s *StoreInfo) IsEvictedAsSlowTrend() bool {
 	return s.slowTrendEvicted.Load() > 0
@@ -295,6 +300,13 @@ func (s *StoreInfo) IsSlow() bool {
 	return s.IsEvictedAsSlowTrend() || s.rawStats.GetSlowScore() >= slowStoreThreshold
 }
 
+// IsStopping checks if the store is in stopping state.
+func (s *StoreInfo) IsStopping() bool {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.rawStats.GetIsStopping()
+}
+
 // GetSlowTrend returns the slow trend information of the store.
 func (s *StoreInfo) GetSlowTrend() *pdpb.SlowTrend {
 	s.mu.RLock()
@@ -946,6 +958,19 @@ func (s *StoresInfo) SlowStoreEvicted(storeID uint64) error {
 	return nil
 }
 
+// StoppingStoreEvicted marks a store as a stopping store and prevents transferring
+// leader to the store
+func (s *StoresInfo) StoppingStoreEvicted(storeID uint64) error {
+	s.Lock()
+	defer s.Unlock()
+	store, ok := s.stores[storeID]
+	if !ok {
+		return errs.ErrStoreNotFound.FastGenByArgs(storeID)
+	}
+	s.stores[storeID] = store.Clone(StoppingStoreEvicted())
+	return nil
+}
+
 // SlowStoreRecovered cleans the evicted state of a store.
 func (s *StoresInfo) SlowStoreRecovered(storeID uint64) {
 	s.Lock()
@@ -959,6 +984,19 @@ func (s *StoresInfo) SlowStoreRecovered(storeID uint64) {
 	s.stores[storeID] = store.Clone(SlowStoreRecovered())
 }
 
+// StoppingStoreRecovered cleans the evicted state of a store.
+func (s *StoresInfo) StoppingStoreRecovered(storeID uint64) {
+	s.Lock()
+	defer s.Unlock()
+	store, ok := s.stores[storeID]
+	if !ok {
+		log.Warn("try to clean a store's evicted as a stopping store state, but it is not found. It may be cleanup",
+			zap.Uint64("store-id", storeID))
+		return
+	}
+	s.stores[storeID] = store.Clone(StoppingStoreRecovered())
+}
+
 // SlowTrendEvicted marks a store as a slow trend and prevents transferring
 // leader to the store
 func (s *StoresInfo) SlowTrendEvicted(storeID uint64) error {

pkg/core/store_option.go

@@ -142,6 +142,14 @@ func SlowStoreEvicted() StoreCreateOption {
 	}
 }
 
+// StoppingStoreEvicted marks a store as a stopping store and prevents transferring
+// leader to the store
+func StoppingStoreEvicted() StoreCreateOption {
+	return func(store *StoreInfo) {
+		store.rawStats.IsStopping = true
+	}
+}
+
 // SlowTrendEvicted marks a store as a slow store by trend and prevents transferring
 // leader to the store
 func SlowTrendEvicted() StoreCreateOption {
@@ -164,6 +172,13 @@ func SlowStoreRecovered() StoreCreateOption {
 	}
 }
 
+// StoppingStoreRecovered cleans the evicted state of a store.
+func StoppingStoreRecovered() StoreCreateOption {
+	return func(store *StoreInfo) {
+		store.rawStats.IsStopping = false
+	}
+}
+
 // SetLeaderCount sets the leader count for the store.
 func SetLeaderCount(leaderCount int) StoreCreateOption {
 	return func(store *StoreInfo) {

pkg/schedule/config/config.go

@@ -590,6 +590,7 @@ var defaultSchedulersInit = func() SchedulerConfigs {
 	}
 	if !kerneltype.IsNextGen() {
 		defaultSchedulers = append(defaultSchedulers, SchedulerConfig{Type: types.SchedulerTypeCompatibleMap[types.EvictSlowStoreScheduler]})
+		defaultSchedulers = append(defaultSchedulers, SchedulerConfig{Type: types.SchedulerTypeCompatibleMap[types.EvictStoppingStoreScheduler]})
 	}
 	return defaultSchedulers
 }

pkg/schedule/filter/counter.go

@@ -59,6 +59,7 @@ const (
 	storeStateOffline
 	storeStatePauseLeader
 	storeStateSlow
+	storeStateStopping
 	storeStateDisconnected
 	storeStateBusy
 	storeStateExceedRemoveLimit
@@ -88,6 +89,7 @@ var filters = [filtersLen]string{
 	"store-state-offline-filter",
 	"store-state-pause-leader-filter",
 	"store-state-slow-filter",
+	"store-state-stopping-filter",
 	"store-state-disconnect-filter",
 	"store-state-busy-filter",
 	"store-state-exceed-remove-limit-filter",

pkg/schedule/filter/filters.go

@@ -390,6 +390,15 @@ func (f *StoreStateFilter) slowStoreEvicted(_ config.SharedConfigProvider, store
 	return statusOK
 }
 
+func (f *StoreStateFilter) stoppingStoreEvicted(_ config.SharedConfigProvider, store *core.StoreInfo) *plan.Status {
+	if store.EvictedAsStoppingStore() {
+		f.Reason = storeStateStopping
+		return statusStoreRejectLeader
+	}
+	f.Reason = storeStateOK
+	return statusOK
+}
+
 func (f *StoreStateFilter) slowTrendEvicted(_ config.SharedConfigProvider, store *core.StoreInfo) *plan.Status {
 	if store.IsEvictedAsSlowTrend() {
 		f.Reason = storeStateSlowTrend
@@ -500,7 +509,7 @@ func (f *StoreStateFilter) anyConditionMatch(typ int, conf config.SharedConfigPr
 		funcs = []conditionFunc{f.isBusy}
 	case leaderTarget:
 		funcs = []conditionFunc{f.isRemoved, f.isRemoving, f.isDown, f.pauseLeaderTransferIn,
-			f.slowStoreEvicted, f.slowTrendEvicted, f.isDisconnected, f.isBusy, f.hasRejectLeaderProperty}
+			f.slowStoreEvicted, f.stoppingStoreEvicted, f.slowTrendEvicted, f.isDisconnected, f.isBusy, f.hasRejectLeaderProperty}
 	case regionTarget:
 		funcs = []conditionFunc{f.isRemoved, f.isRemoving, f.isDown, f.isDisconnected, f.isBusy,
 			f.exceedAddLimit, f.tooManySnapshots, f.tooManyPendingPeers}