fix

Signed-off-by: husharp <jinhao.hu@pingcap.com>

Author: HuSharp

pkg/core/region.go

@@ -71,8 +71,14 @@ type RegionInfo struct {
 	queryStats        *pdpb.QueryStats
 	flowRoundDivisor  uint64
 	// buckets is not thread unsafe, it should be accessed by the request `report buckets` with greater version.
-	buckets       unsafe.Pointer
-	fromHeartbeat bool
+	buckets unsafe.Pointer
+	// source is used to indicate region's source, such as FromHeartbeat/FromSync/InDisk.
+	source RegionSource
+}
+
+// GetRegionSource returns the region source.
+func (r *RegionInfo) GetRegionSource() RegionSource {
+	return r.source
 }
 
 // NewRegionInfo creates RegionInfo with region's meta and leader peer.
@@ -171,6 +177,7 @@ func RegionFromHeartbeat(heartbeat *pdpb.RegionHeartbeatRequest, opts ...RegionC
 		interval:          heartbeat.GetInterval(),
 		replicationStatus: heartbeat.GetReplicationStatus(),
 		queryStats:        heartbeat.GetQueryStats(),
+		source:            FromHeartbeat,
 	}
 
 	for _, opt := range opts {
@@ -639,11 +646,6 @@ func (r *RegionInfo) IsFlashbackChanged(l *RegionInfo) bool {
 	return r.meta.FlashbackStartTs != l.meta.FlashbackStartTs || r.meta.IsInFlashback != l.meta.IsInFlashback
 }
 
-// IsFromHeartbeat returns whether the region info is from the region heartbeat.
-func (r *RegionInfo) IsFromHeartbeat() bool {
-	return r.fromHeartbeat
-}
-
 func (r *RegionInfo) isInvolved(startKey, endKey []byte) bool {
 	return bytes.Compare(r.GetStartKey(), startKey) >= 0 && (len(endKey) == 0 || (len(r.GetEndKey()) > 0 && bytes.Compare(r.GetEndKey(), endKey) <= 0))
 }
@@ -683,7 +685,7 @@ func GenerateRegionGuideFunc(enableLog bool) RegionGuideFunc {
 			}
 			saveKV, saveCache, isNew = true, true, true
 		} else {
-			if !origin.IsFromHeartbeat() {
+			if origin.source == FromSync || origin.source == InDisk {
 				isNew = true
 			}
 			r := region.GetRegionEpoch()
@@ -793,6 +795,18 @@ type RegionsInfo struct {
 	pendingPeers map[uint64]*regionTree // storeID -> sub regionTree
 }
 
+// RegionSource is the source of region.
+type RegionSource uint32
+
+const (
+	// InDisk means region is stale.
+	InDisk RegionSource = iota
+	// FromSync means region is stale.
+	FromSync
+	// FromHeartbeat means region is fresh.
+	FromHeartbeat
+)
+
 // NewRegionsInfo creates RegionsInfo with tree, regions, leaders and followers
 func NewRegionsInfo() *RegionsInfo {
 	return &RegionsInfo{
@@ -840,6 +854,8 @@ func (r *RegionsInfo) CheckAndPutRegion(region *RegionInfo) []*RegionInfo {
 	origin, overlaps, rangeChanged := r.setRegionLocked(region, true, ols...)
 	r.t.Unlock()
 	r.UpdateSubTree(region, origin, overlaps, rangeChanged)
+	// InDisk means region is stale.
+	r.AtomicAddStaleRegionCnt()
 	return overlaps
 }
 
@@ -857,6 +873,21 @@ func (r *RegionsInfo) PreCheckPutRegion(region *RegionInfo) (*RegionInfo, []*reg
 	return origin, overlaps, err
 }
 
+// GetStaleRegionCnt returns the stale region count.
+func (r *RegionsInfo) GetStaleRegionCnt() int64 {
+	r.t.RLock()
+	defer r.t.RUnlock()
+	if r.tree.length() == 0 {
+		return 0
+	}
+	return r.tree.staleRegionCnt
+}
+
+// AtomicAddStaleRegionCnt atomically adds the stale region count.
+func (r *RegionsInfo) AtomicAddStaleRegionCnt() {
+	r.tree.AtomicAddStaleRegionCnt()
+}
+
 // AtomicCheckAndPutRegion checks if the region is valid to put, if valid then put.
 func (r *RegionsInfo) AtomicCheckAndPutRegion(region *RegionInfo) ([]*RegionInfo, error) {
 	r.t.Lock()
@@ -870,6 +901,10 @@ func (r *RegionsInfo) AtomicCheckAndPutRegion(region *RegionInfo) ([]*RegionInfo
 		r.t.Unlock()
 		return nil, err
 	}
+	// If origin is stale, need to sub the stale region count.
+	if origin != nil && origin.source != FromHeartbeat && region.source == FromHeartbeat {
+		r.tree.AtomicSubStaleRegionCnt()
+	}
 	origin, overlaps, rangeChanged := r.setRegionLocked(region, true, ols...)
 	r.t.Unlock()
 	r.UpdateSubTree(region, origin, overlaps, rangeChanged)

pkg/core/region_option.go

@@ -381,10 +381,10 @@ func WithInterval(interval *pdpb.TimeInterval) RegionCreateOption {
 	}
 }
 
-// SetFromHeartbeat sets if the region info comes from the region heartbeat.
-func SetFromHeartbeat(fromHeartbeat bool) RegionCreateOption {
+// SetSource sets the region info's come from.
+func SetSource(source RegionSource) RegionCreateOption {
 	return func(region *RegionInfo) {
-		region.fromHeartbeat = fromHeartbeat
+		region.source = source
 	}
 }
 

pkg/core/region_tree.go

@@ -16,6 +16,7 @@ package core
 import (
 	"bytes"
 	"math/rand"
+	"sync/atomic"
 
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/log"
@@ -61,6 +62,8 @@ type regionTree struct {
 	totalSize           int64
 	totalWriteBytesRate float64
 	totalWriteKeysRate  float64
+
+	staleRegionCnt int64
 }
 
 func newRegionTree() *regionTree {
@@ -69,6 +72,7 @@ func newRegionTree() *regionTree {
 		totalSize:           0,
 		totalWriteBytesRate: 0,
 		totalWriteKeysRate:  0,
+		staleRegionCnt:      0,
 	}
 }
 
@@ -342,3 +346,20 @@ func (t *regionTree) TotalWriteRate() (bytesRate, keysRate float64) {
 	}
 	return t.totalWriteBytesRate, t.totalWriteKeysRate
 }
+
+func (t *regionTree) AtomicAddStaleRegionCnt() {
+	if t.length() == 0 {
+		return
+	}
+	atomic.AddInt64(&t.staleRegionCnt, 1)
+}
+
+func (t *regionTree) AtomicSubStaleRegionCnt() {
+	if t.length() == 0 {
+		return
+	}
+	if t.staleRegionCnt == 0 {
+		return
+	}
+	atomic.AddInt64(&t.staleRegionCnt, -1)
+}

pkg/schedule/prepare_checker.go

@@ -17,8 +17,10 @@ package schedule
 import (
 	"time"
 
+	"github.com/pingcap/log"
 	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/utils/syncutil"
+	"go.uber.org/zap"
 )
 
 type prepareChecker struct {
@@ -47,6 +49,11 @@ func (checker *prepareChecker) check(c *core.BasicCluster) bool {
 		checker.prepared = true
 		return true
 	}
+	if float64(c.GetStaleRegionCnt()) < float64(c.GetTotalRegionCount())*(1-collectFactor) {
+		log.Info("stale region num is satisfied, skip prepare checker", zap.Int64("stale-region", c.GetStaleRegionCnt()), zap.Int("total-region", c.GetTotalRegionCount()))
+		checker.prepared = true
+		return true
+	}
 	// The number of active regions should be more than total region of all stores * collectFactor
 	if float64(c.GetTotalRegionCount())*collectFactor > float64(checker.sum) {
 		return false

pkg/syncer/client.go

@@ -183,17 +183,21 @@ func (s *RegionSyncer) StartSyncWithLeader(addr string) {
 							core.SetWrittenKeys(stats[i].KeysWritten),
 							core.SetReadBytes(stats[i].BytesRead),
 							core.SetReadKeys(stats[i].KeysRead),
-							core.SetFromHeartbeat(false),
+							core.SetSource(core.FromSync),
 						)
 					} else {
-						region = core.NewRegionInfo(r, regionLeader, core.SetFromHeartbeat(false))
+						region = core.NewRegionInfo(r, regionLeader, core.SetSource(core.FromSync))
 					}
 
 					origin, _, err := bc.PreCheckPutRegion(region)
 					if err != nil {
 						log.Debug("region is stale", zap.Stringer("origin", origin.GetMeta()), errs.ZapError(err))
 						continue
 					}
+					// FromSync means region is stale.
+					if origin == nil || (origin != nil && origin.GetRegionSource() == core.FromHeartbeat) {
+						bc.RegionsInfo.AtomicAddStaleRegionCnt()
+					}
 					_, saveKV, _, _ := regionGuide(region, origin)
 					overlaps := bc.PutRegion(region)
 

server/cluster/cluster.go

@@ -1128,7 +1128,7 @@ func (c *RaftCluster) processRegionHeartbeat(region *core.RegionInfo) error {
 		// To prevent a concurrent heartbeat of another region from overriding the up-to-date region info by a stale one,
 		// check its validation again here.
 		//
-		// However it can't solve the race condition of concurrent heartbeats from the same region.
+		// However, it can't solve the race condition of concurrent heartbeats from the same region.
 		if overlaps, err = c.core.AtomicCheckAndPutRegion(region); err != nil {
 			return err
 		}

server/cluster/cluster_test.go

@@ -1002,7 +1002,7 @@ func TestRegionSizeChanged(t *testing.T) {
 		core.WithLeader(region.GetPeers()[2]),
 		core.SetApproximateSize(curMaxMergeSize-1),
 		core.SetApproximateKeys(curMaxMergeKeys-1),
-		core.SetFromHeartbeat(true),
+		core.SetSource(core.FromHeartbeat),
 	)
 	cluster.processRegionHeartbeat(region)
 	regionID := region.GetID()
@@ -1012,7 +1012,7 @@ func TestRegionSizeChanged(t *testing.T) {
 		core.WithLeader(region.GetPeers()[2]),
 		core.SetApproximateSize(curMaxMergeSize+1),
 		core.SetApproximateKeys(curMaxMergeKeys+1),
-		core.SetFromHeartbeat(true),
+		core.SetSource(core.FromHeartbeat),
 	)
 	cluster.processRegionHeartbeat(region)
 	re.False(cluster.regionStats.IsRegionStatsType(regionID, statistics.UndersizedRegion))
@@ -2375,6 +2375,7 @@ func (c *testCluster) LoadRegion(regionID uint64, followerStoreIDs ...uint64) er
 		peer, _ := c.AllocPeer(id)
 		region.Peers = append(region.Peers, peer)
 	}
+	c.core.AtomicAddStaleRegionCnt()
 	return c.putRegion(core.NewRegionInfo(region, nil))
 }
 

server/grpc_service.go

@@ -1261,7 +1261,7 @@ func (s *GrpcServer) RegionHeartbeat(stream pdpb.PD_RegionHeartbeatServer) error
 			lastBind = time.Now()
 		}
 
-		region := core.RegionFromHeartbeat(request, flowRoundOption, core.SetFromHeartbeat(true))
+		region := core.RegionFromHeartbeat(request, flowRoundOption)
 		if region.GetLeader() == nil {
 			log.Error("invalid request, the leader is nil", zap.Reflect("request", request), errs.ZapError(errs.ErrLeaderNil))
 			regionHeartbeatCounter.WithLabelValues(storeAddress, storeLabel, "report", "invalid-leader").Inc()