core/region: fix the bug causing incorrect reference count after a ca…

…che reset (#8201)

ref #7897

#8132 introduced a count reference to ensure the consistency between
the root tree and subtrees, however, the reference count was not handled correctly during the cache
reset process, which may lead to incorrect reference counts in later heartbeat processing.

Signed-off-by: JmPotato <[email protected]>

pkg/core/region.go

@@ -929,7 +929,7 @@ func NewRegionsInfo() *RegionsInfo {
 		learners:     make(map[uint64]*regionTree),
 		witnesses:    make(map[uint64]*regionTree),
 		pendingPeers: make(map[uint64]*regionTree),
-		overlapTree:  newRegionTree(),
+		overlapTree:  newRegionTreeWithCountRef(),
 	}
 }
 
@@ -1131,42 +1131,31 @@ func (r *RegionsInfo) updateSubTreeLocked(rangeChanged bool, overlaps []*RegionI
 	item := &regionItem{region}
 	r.subRegions[region.GetID()] = item
 	r.overlapTree.update(item, false)
-	setPeer := func(peersMap map[uint64]*regionTree, storeID uint64, item *regionItem, countRef bool) {
+	// Add leaders and followers.
+	setPeer := func(peersMap map[uint64]*regionTree, storeID uint64) {
 		store, ok := peersMap[storeID]
 		if !ok {
-			if !countRef {
-				store = newRegionTree()
-			} else {
-				store = newRegionTreeWithCountRef()
-			}
+			store = newRegionTree()
 			peersMap[storeID] = store
 		}
 		store.update(item, false)
 	}
-
-	// Add to leaders and followers.
 	for _, peer := range region.GetVoters() {
 		storeID := peer.GetStoreId()
 		if peer.GetId() == region.leader.GetId() {
-			// Add leader peer to leaders.
-			setPeer(r.leaders, storeID, item, true)
+			setPeer(r.leaders, storeID)
 		} else {
-			// Add follower peer to followers.
-			setPeer(r.followers, storeID, item, false)
+			setPeer(r.followers, storeID)
 		}
 	}
-
+	// Add other peers.
 	setPeers := func(peersMap map[uint64]*regionTree, peers []*metapb.Peer) {
 		for _, peer := range peers {
-			storeID := peer.GetStoreId()
-			setPeer(peersMap, storeID, item, false)
+			setPeer(peersMap, peer.GetStoreId())
 		}
 	}
-	// Add to learners.
 	setPeers(r.learners, region.GetLearners())
-	// Add to witnesses.
 	setPeers(r.witnesses, region.GetWitnesses())
-	// Add to PendingPeers
 	setPeers(r.pendingPeers, region.GetPendingPeers())
 }
 
@@ -1335,7 +1324,7 @@ func (r *RegionsInfo) RemoveRegion(region *RegionInfo) {
 // ResetRegionCache resets the regions info.
 func (r *RegionsInfo) ResetRegionCache() {
 	r.t.Lock()
-	r.tree = newRegionTree()
+	r.tree = newRegionTreeWithCountRef()
 	r.regions = make(map[uint64]*regionItem)
 	r.t.Unlock()
 	r.st.Lock()
@@ -1345,7 +1334,7 @@ func (r *RegionsInfo) ResetRegionCache() {
 	r.learners = make(map[uint64]*regionTree)
 	r.witnesses = make(map[uint64]*regionTree)
 	r.pendingPeers = make(map[uint64]*regionTree)
-	r.overlapTree = newRegionTree()
+	r.overlapTree = newRegionTreeWithCountRef()
 }
 
 // RemoveRegionFromSubTree removes RegionInfo from regionSubTrees

pkg/core/region_test.go

@@ -1081,3 +1081,18 @@ func TestCheckAndPutSubTree(t *testing.T) {
 	// should failed to put because the root tree is missing
 	re.Equal(0, regions.tree.length())
 }
+
+func TestCntRefAfterResetRegionCache(t *testing.T) {
+	re := require.New(t)
+	regions := NewRegionsInfo()
+	// Put the region first.
+	region := NewTestRegionInfo(1, 1, []byte("a"), []byte("b"))
+	regions.CheckAndPutRegion(region)
+	re.Equal(int32(2), region.GetRef())
+	regions.ResetRegionCache()
+	// Put the region after reset.
+	region = NewTestRegionInfo(1, 1, []byte("a"), []byte("b"))
+	re.Zero(region.GetRef())
+	regions.CheckAndPutRegion(region)
+	re.Equal(int32(2), region.GetRef())
+}