From b414fcf67ad77050b4e5a6fd54a182e9f9e16758 Mon Sep 17 00:00:00 2001
From: disksing <i@disksing.com>
Date: Thu, 18 Nov 2021 11:34:05 +0800
Subject: [PATCH] *: check if GetStore returns nil

close #4344

Signed-off-by: disksing <i@disksing.com>
---
 server/schedule/checker/replica_strategy.go | 6 +++++-
 server/schedule/region_scatterer.go         | 9 +++++++++
 server/schedulers/shuffle_region.go         | 6 +++++-
 server/statistics/store.go                  | 2 +-
 4 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/server/schedule/checker/replica_strategy.go b/server/schedule/checker/replica_strategy.go
index 1adbf383ad6..0f816df28c6 100644
--- a/server/schedule/checker/replica_strategy.go
+++ b/server/schedule/checker/replica_strategy.go
@@ -93,8 +93,12 @@ func (s *ReplicaStrategy) SelectStoreToFix(coLocationStores []*core.StoreInfo, o
 func (s *ReplicaStrategy) SelectStoreToImprove(coLocationStores []*core.StoreInfo, old uint64) uint64 {
 	// trick to avoid creating a slice with `old` removed.
 	s.swapStoreToFirst(coLocationStores, old)
+	oldStore := s.cluster.GetStore(old)
+	if oldStore == nil {
+		return 0
+	}
 	filters := []filter.Filter{
-		filter.NewLocationImprover(s.checkerName, s.locationLabels, coLocationStores, s.cluster.GetStore(old)),
+		filter.NewLocationImprover(s.checkerName, s.locationLabels, coLocationStores, oldStore),
 	}
 	if len(s.locationLabels) > 0 && s.isolationLevel != "" {
 		filters = append(filters, filter.NewIsolationFilter(s.checkerName, s.isolationLevel, s.locationLabels, coLocationStores[1:]))
diff --git a/server/schedule/region_scatterer.go b/server/schedule/region_scatterer.go
index 87e9193dbf7..de8bbecffc1 100644
--- a/server/schedule/region_scatterer.go
+++ b/server/schedule/region_scatterer.go
@@ -278,6 +278,9 @@ func (r *RegionScatterer) scatterRegion(region *core.RegionInfo, group string) *
 	// Group peers by the engine of their stores
 	for _, peer := range region.GetPeers() {
 		store := r.cluster.GetStore(peer.GetStoreId())
+		if store == nil {
+			return nil
+		}
 		if ordinaryFilter.Target(r.cluster.GetOpts(), store) {
 			ordinaryPeers[peer.GetId()] = peer
 		} else {
@@ -425,6 +428,9 @@ func (r *RegionScatterer) selectAvailableLeaderStore(group string, peers map[uin
 	leaderCandidateStores := make([]uint64, 0)
 	for storeID := range peers {
 		store := r.cluster.GetStore(storeID)
+		if store == nil {
+			return 0
+		}
 		engine := store.GetLabelValue(core.EngineKey)
 		if len(engine) < 1 {
 			leaderCandidateStores = append(leaderCandidateStores, storeID)
@@ -449,6 +455,9 @@ func (r *RegionScatterer) Put(peers map[uint64]*metapb.Peer, leaderStoreID uint6
 	for _, peer := range peers {
 		storeID := peer.GetStoreId()
 		store := r.cluster.GetStore(storeID)
+		if store == nil {
+			continue
+		}
 		if ordinaryFilter.Target(r.cluster.GetOpts(), store) {
 			r.ordinaryEngine.selectedPeer.Put(storeID, group)
 			scatterDistributionCounter.WithLabelValues(
diff --git a/server/schedulers/shuffle_region.go b/server/schedulers/shuffle_region.go
index 487e3bcaebb..996dc2416a6 100644
--- a/server/schedulers/shuffle_region.go
+++ b/server/schedulers/shuffle_region.go
@@ -153,7 +153,11 @@ func (s *shuffleRegionScheduler) scheduleRemovePeer(cluster opt.Cluster) (*core.
 }
 
 func (s *shuffleRegionScheduler) scheduleAddPeer(cluster opt.Cluster, region *core.RegionInfo, oldPeer *metapb.Peer) *metapb.Peer {
-	scoreGuard := filter.NewPlacementSafeguard(s.GetName(), cluster, region, cluster.GetStore(oldPeer.GetStoreId()))
+	store := cluster.GetStore(oldPeer.GetStoreId())
+	if store == nil {
+		return nil
+	}
+	scoreGuard := filter.NewPlacementSafeguard(s.GetName(), cluster, region, store)
 	excludedFilter := filter.NewExcludedFilter(s.GetName(), nil, region.GetStoreIds())
 
 	target := filter.NewCandidates(cluster.GetStores()).
diff --git a/server/statistics/store.go b/server/statistics/store.go
index e8cb5ae139f..a10234696f5 100644
--- a/server/statistics/store.go
+++ b/server/statistics/store.go
@@ -119,7 +119,7 @@ func (s *StoresStats) FilterUnhealthyStore(cluster core.StoreSetInformer) {
 	defer s.Unlock()
 	for storeID := range s.rollingStoresStats {
 		store := cluster.GetStore(storeID)
-		if store.IsTombstone() || store.IsUnhealthy() || store.IsPhysicallyDestroyed() {
+		if store == nil || store.IsTombstone() || store.IsUnhealthy() || store.IsPhysicallyDestroyed() {
 			delete(s.rollingStoresStats, storeID)
 		}
 	}