From 19eda096c99b0e8cd51188eab0e3351227dd058d Mon Sep 17 00:00:00 2001
From: JmPotato <ghzpotato@gmail.com>
Date: Tue, 21 May 2024 09:47:45 +0800
Subject: [PATCH] core/region: fix the bug causing incorrect reference count
 after a cache reset (#8201)

ref tikv/pd#7897

https://github.com/tikv/pd/pull/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 <ghzpotato@gmail.com>
---
 pkg/core/region.go      | 31 ++++++++++---------------------
 pkg/core/region_test.go | 15 +++++++++++++++
 2 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/pkg/core/region.go b/pkg/core/region.go
index 54091968d1a..c9a8455d4de 100644
--- a/pkg/core/region.go
+++ b/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
diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go
index b15e8f38b93..1b8f20cf9b2 100644
--- a/pkg/core/region_test.go
+++ b/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())
+}