*: reduce get overlap operations

server/cluster/cluster.go

@@ -829,7 +829,7 @@ var regionGuide = core.GenerateRegionGuideFunc(true)
 
 // processRegionHeartbeat updates the region information.
 func (c *RaftCluster) processRegionHeartbeat(region *core.RegionInfo) error {
-	origin, err := c.core.PreCheckPutRegion(region)
+	origin, _, err := c.core.PreCheckPutRegion(region)
 	if err != nil {
 		return err
 	}

server/cluster/cluster_test.go

@@ -1678,7 +1678,7 @@ func Test(t *testing.T) {
 		re.Nil(cache.GetRegionByKey(regionKey))
 		checkRegions(re, cache, regions[0:i])
 
-		origin, overlaps, rangeChanged := cache.SetRegionWithUpdate(region)
+		origin, overlaps, rangeChanged := cache.SetRegion(region)
 		cache.UpdateSubTree(region, origin, overlaps, rangeChanged)
 		checkRegion(re, cache.GetRegion(i), region)
 		checkRegion(re, cache.GetRegionByKey(regionKey), region)
@@ -1692,7 +1692,7 @@ func Test(t *testing.T) {
 		// Update leader to peer np-1.
 		newRegion := region.Clone(core.WithLeader(region.GetPeers()[np-1]))
 		regions[i] = newRegion
-		origin, overlaps, rangeChanged = cache.SetRegionWithUpdate(newRegion)
+		origin, overlaps, rangeChanged = cache.SetRegion(newRegion)
 		cache.UpdateSubTree(newRegion, origin, overlaps, rangeChanged)
 		checkRegion(re, cache.GetRegion(i), newRegion)
 		checkRegion(re, cache.GetRegionByKey(regionKey), newRegion)
@@ -1707,7 +1707,7 @@ func Test(t *testing.T) {
 		// Reset leader to peer 0.
 		newRegion = region.Clone(core.WithLeader(region.GetPeers()[0]))
 		regions[i] = newRegion
-		origin, overlaps, rangeChanged = cache.SetRegionWithUpdate(newRegion)
+		origin, overlaps, rangeChanged = cache.SetRegion(newRegion)
 		cache.UpdateSubTree(newRegion, origin, overlaps, rangeChanged)
 		checkRegion(re, cache.GetRegion(i), newRegion)
 		checkRegions(re, cache, regions[0:(i+1)])
@@ -1729,7 +1729,7 @@ func Test(t *testing.T) {
 	// check overlaps
 	// clone it otherwise there are two items with the same key in the tree
 	overlapRegion := regions[n-1].Clone(core.WithStartKey(regions[n-2].GetStartKey()))
-	origin, overlaps, rangeChanged := cache.SetRegionWithUpdate(overlapRegion)
+	origin, overlaps, rangeChanged := cache.SetRegion(overlapRegion)
 	cache.UpdateSubTree(overlapRegion, origin, overlaps, rangeChanged)
 	re.Nil(cache.GetRegion(n - 2))
 	re.NotNil(cache.GetRegion(n - 1))
@@ -1739,7 +1739,7 @@ func Test(t *testing.T) {
 		for j := 0; j < cache.GetStoreLeaderCount(i); j++ {
 			region := filter.SelectOneRegion(tc.RandLeaderRegions(i, []core.KeyRange{core.NewKeyRange("", "")}), nil, pendingFilter, downFilter)
 			newRegion := region.Clone(core.WithPendingPeers(region.GetPeers()))
-			origin, overlaps, rangeChanged = cache.SetRegionWithUpdate(newRegion)
+			origin, overlaps, rangeChanged = cache.SetRegion(newRegion)
 			cache.UpdateSubTree(newRegion, origin, overlaps, rangeChanged)
 		}
 		re.Nil(filter.SelectOneRegion(tc.RandLeaderRegions(i, []core.KeyRange{core.NewKeyRange("", "")}), nil, pendingFilter, downFilter))

server/core/basic_cluster.go

@@ -16,11 +16,8 @@ package core
 
 import (
 	"github.com/pingcap/kvproto/pkg/metapb"
-	"github.com/pingcap/log"
-	"github.com/tikv/pd/pkg/errs"
 	"github.com/tikv/pd/pkg/syncutil"
 	"github.com/tikv/pd/server/core/storelimit"
-	"go.uber.org/zap"
 )
 
 // BasicCluster provides basic data member and interface for a tikv cluster.
@@ -221,45 +218,6 @@ func (bc *BasicCluster) GetStoresWriteRate() (storeIDs []uint64, bytesRates, key
 	return bc.getWriteRate(bc.RegionsInfo.GetStoreWriteRate)
 }
 
-/* Regions write operations */
-
-// PreCheckPutRegion checks if the region is valid to put.
-func (bc *BasicCluster) PreCheckPutRegion(region *RegionInfo) (*RegionInfo, error) {
-	origin, overlaps := bc.RegionsInfo.GetRelevantRegions(region)
-	return check(region, origin, overlaps)
-}
-
-// CheckAndPutRegion checks if the region is valid to put, if valid then put.
-func (bc *BasicCluster) CheckAndPutRegion(region *RegionInfo) []*RegionInfo {
-	origin, err := bc.PreCheckPutRegion(region)
-	if err != nil {
-		log.Debug("region is stale", zap.Stringer("origin", origin.GetMeta()), errs.ZapError(err))
-		// return the state region to delete.
-		return []*RegionInfo{region}
-	}
-	return bc.PutRegion(region)
-}
-
-// PutRegion put a region.
-func (bc *BasicCluster) PutRegion(region *RegionInfo) []*RegionInfo {
-	origin, overlaps, rangeChanged := bc.RegionsInfo.SetRegionWithUpdate(region)
-	bc.RegionsInfo.UpdateSubTree(region, origin, overlaps, rangeChanged)
-	return overlaps
-}
-
-// RemoveRegionIfExist removes RegionInfo from regionTree and regionMap if exists.
-func (bc *BasicCluster) RemoveRegionIfExist(id uint64) {
-	if r := bc.RegionsInfo.GetRegion(id); r != nil {
-		bc.RegionsInfo.RemoveRegion(r)
-		bc.RegionsInfo.RemoveRegionFromSubTree(r)
-	}
-}
-
-// ResetRegionCache drops all region cache.
-func (bc *BasicCluster) ResetRegionCache() {
-	bc.RegionsInfo.Reset()
-}
-
 // RegionSetInformer provides access to a shared informer of regions.
 type RegionSetInformer interface {
 	GetRegionCount() int

server/core/region.go

@@ -31,6 +31,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/kvproto/pkg/replication_modepb"
 	"github.com/pingcap/log"
+	"github.com/tikv/pd/pkg/errs"
 	"github.com/tikv/pd/pkg/logutil"
 	"github.com/tikv/pd/pkg/syncutil"
 	"github.com/tikv/pd/pkg/typeutil"
@@ -746,45 +747,79 @@ func (r *RegionsInfo) getRegionLocked(regionID uint64) *RegionInfo {
 	return nil
 }
 
+// CheckAndPutRegion checks if the region is valid to put, if valid then put.
+func (r *RegionsInfo) CheckAndPutRegion(region *RegionInfo) []*RegionInfo {
+	r.t.Lock()
+	origin := r.getRegionLocked(region.GetID())
+	var ols []*regionItem
+	if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) {
+		ols = r.tree.overlaps(&regionItem{RegionInfo: region})
+	}
+	err := check(region, origin, ols)
+	if err != nil {
+		log.Debug("region is stale", zap.Stringer("origin", origin.GetMeta()), errs.ZapError(err))
+		// return the state region to delete.
+		return []*RegionInfo{region}
+	}
+	origin, overlaps, rangeChanged := r.setRegionLocked(region, true, ols...)
+	r.t.Unlock()
+	r.UpdateSubTree(region, origin, overlaps, rangeChanged)
+	return overlaps
+}
+
+// PutRegion put a region.
+func (r *RegionsInfo) PutRegion(region *RegionInfo) []*RegionInfo {
+	origin, overlaps, rangeChanged := r.SetRegion(region)
+	r.UpdateSubTree(region, origin, overlaps, rangeChanged)
+	return overlaps
+}
+
+// PreCheckPutRegion checks if the region is valid to put.
+func (r *RegionsInfo) PreCheckPutRegion(region *RegionInfo) (*RegionInfo, []*regionItem, error) {
+	origin, overlaps := r.GetRelevantRegions(region)
+	err := check(region, origin, overlaps)
+	return origin, overlaps, err
+}
+
 // AtomicCheckAndPutRegion checks if the region is valid to put, if valid then put.
 func (r *RegionsInfo) AtomicCheckAndPutRegion(region *RegionInfo) ([]*RegionInfo, error) {
 	r.t.Lock()
-	var overlaps []*RegionInfo
+	var ols []*regionItem
 	origin := r.getRegionLocked(region.GetID())
 	if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) {
-		overlaps = r.tree.getOverlaps(region)
+		ols = r.tree.overlaps(&regionItem{RegionInfo: region})
 	}
-	_, err := check(region, origin, overlaps)
+	err := check(region, origin, ols)
 	if err != nil {
 		r.t.Unlock()
 		return nil, err
 	}
-	origin, overlaps, rangeChanged := r.setRegionLocked(region)
+	origin, overlaps, rangeChanged := r.setRegionLocked(region, true, ols...)
 	r.t.Unlock()
 	r.UpdateSubTree(region, origin, overlaps, rangeChanged)
 	return overlaps, nil
 }
 
 // GetRelevantRegions returns the relevant regions for a given region.
-func (r *RegionsInfo) GetRelevantRegions(region *RegionInfo) (origin *RegionInfo, overlaps []*RegionInfo) {
+func (r *RegionsInfo) GetRelevantRegions(region *RegionInfo) (origin *RegionInfo, overlaps []*regionItem) {
 	r.t.RLock()
 	defer r.t.RUnlock()
 	origin = r.getRegionLocked(region.GetID())
 	if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) {
-		overlaps = r.tree.getOverlaps(region)
+		overlaps = r.tree.overlaps(&regionItem{RegionInfo: region})
 	}
 	return
 }
 
-func check(region, origin *RegionInfo, overlaps []*RegionInfo) (*RegionInfo, error) {
+func check(region, origin *RegionInfo, overlaps []*regionItem) error {
 	for _, item := range overlaps {
 		// PD ignores stale regions' heartbeats, unless it is recreated recently by unsafe recover operation.
 		if region.GetRegionEpoch().GetVersion() < item.GetRegionEpoch().GetVersion() && !region.isRegionRecreated() {
-			return nil, errRegionIsStale(region.GetMeta(), item.GetMeta())
+			return errRegionIsStale(region.GetMeta(), item.GetMeta())
 		}
 	}
 	if origin == nil {
-		return nil, nil
+		return nil
 	}
 
 	r := region.GetRegionEpoch()
@@ -793,29 +828,20 @@ func check(region, origin *RegionInfo, overlaps []*RegionInfo) (*RegionInfo, err
 	isTermBehind := region.GetTerm() > 0 && region.GetTerm() < origin.GetTerm()
 	// Region meta is stale, return an error.
 	if (isTermBehind || r.GetVersion() < o.GetVersion() || r.GetConfVer() < o.GetConfVer()) && !region.isRegionRecreated() {
-		return origin, errRegionIsStale(region.GetMeta(), origin.GetMeta())
+		return errRegionIsStale(region.GetMeta(), origin.GetMeta())
 	}
 
-	return origin, nil
-}
-
-// SetRegion sets the RegionInfo to regionTree and regionMap, also update leaders and followers by region peers
-// overlaps: Other regions that overlap with the specified region, excluding itself.
-func (r *RegionsInfo) SetRegion(region *RegionInfo) []*RegionInfo {
-	r.t.Lock()
-	defer r.t.Unlock()
-	_, overlaps, _ := r.setRegionLocked(region)
-	return overlaps
+	return nil
 }
 
-// SetRegionWithUpdate sets the RegionInfo to regionTree and regionMap and return the update info of subtree.
-func (r *RegionsInfo) SetRegionWithUpdate(region *RegionInfo) (*RegionInfo, []*RegionInfo, bool) {
+// SetRegion sets the RegionInfo to regionTree and regionMap and return the update info of subtree.
+func (r *RegionsInfo) SetRegion(region *RegionInfo) (*RegionInfo, []*RegionInfo, bool) {
 	r.t.Lock()
 	defer r.t.Unlock()
-	return r.setRegionLocked(region)
+	return r.setRegionLocked(region, false)
 }
 
-func (r *RegionsInfo) setRegionLocked(region *RegionInfo) (*RegionInfo, []*RegionInfo, bool) {
+func (r *RegionsInfo) setRegionLocked(region *RegionInfo, withOverlaps bool, ol ...*regionItem) (*RegionInfo, []*RegionInfo, bool) {
 	var (
 		item   *regionItem // Pointer to the *RegionInfo of this ID.
 		origin *RegionInfo
@@ -829,6 +855,16 @@ func (r *RegionsInfo) setRegionLocked(region *RegionInfo) (*RegionInfo, []*Regio
 		if rangeChanged {
 			// Delete itself in regionTree so that overlaps will not contain itself.
 			// Because the regionItem is reused, there is no need to delete it in the regionMap.
+			idx := -1
+			for i, o := range ol {
+				if o.GetID() == region.GetID() {
+					idx = i
+					break
+				}
+			}
+			if idx >= 0 {
+				ol = append(ol[:idx], ol[idx+1:]...)
+			}
 			r.tree.remove(origin)
 			// Update the RegionInfo in the regionItem.
 			item.RegionInfo = region
@@ -847,20 +883,17 @@ func (r *RegionsInfo) setRegionLocked(region *RegionInfo) (*RegionInfo, []*Regio
 
 	var overlaps []*RegionInfo
 	if rangeChanged {
-		// It has been removed and all information needs to be updated again.
-		overlaps = r.tree.update(item)
+		overlaps = r.tree.update(item, withOverlaps, ol...)
 		for _, old := range overlaps {
-			o := r.getRegionLocked(old.GetID())
-			// Remove from tree and regions.
-			r.tree.remove(o)
-			delete(r.regions, o.GetID())
+			delete(r.regions, old.GetID())
 		}
 	}
+	// return rangeChanged to prevent duplicated calculation
 	return origin, overlaps, rangeChanged
 }
 
 // UpdateSubTree updates the subtree.
-func (r *RegionsInfo) UpdateSubTree(region, origin *RegionInfo, toRemove []*RegionInfo, rangeChanged bool) {
+func (r *RegionsInfo) UpdateSubTree(region, origin *RegionInfo, overlaps []*RegionInfo, rangeChanged bool) {
 	r.st.Lock()
 	defer r.st.Unlock()
 	if origin != nil {
@@ -875,7 +908,7 @@ func (r *RegionsInfo) UpdateSubTree(region, origin *RegionInfo, toRemove []*Regi
 		}
 	}
 	if rangeChanged {
-		for _, re := range toRemove {
+		for _, re := range overlaps {
 			r.removeRegionFromSubTreeLocked(re)
 		}
 	}
@@ -890,7 +923,7 @@ func (r *RegionsInfo) UpdateSubTree(region, origin *RegionInfo, toRemove []*Regi
 			store = newRegionTree()
 			peersMap[storeID] = store
 		}
-		store.update(item)
+		store.update(item, false)
 	}
 
 	// Add to leaders and followers.
@@ -952,10 +985,10 @@ func (r *RegionsInfo) TreeLen() int {
 }
 
 // GetOverlaps returns the regions which are overlapped with the specified region range.
-func (r *RegionsInfo) GetOverlaps(region *RegionInfo) []*RegionInfo {
+func (r *RegionsInfo) GetOverlaps(region *RegionInfo) []*regionItem {
 	r.t.RLock()
 	defer r.t.RUnlock()
-	return r.tree.getOverlaps(region)
+	return r.tree.overlaps(&regionItem{RegionInfo: region})
 }
 
 // RemoveRegion removes RegionInfo from regionTree and regionMap
@@ -967,8 +1000,8 @@ func (r *RegionsInfo) RemoveRegion(region *RegionInfo) {
 	delete(r.regions, region.GetID())
 }
 
-// Reset resets the regions info.
-func (r *RegionsInfo) Reset() {
+// ResetRegionCache resets the regions info.
+func (r *RegionsInfo) ResetRegionCache() {
 	r.t.Lock()
 	r.tree = newRegionTree()
 	r.regions = make(map[uint64]*regionItem)
@@ -1004,6 +1037,14 @@ func (r *RegionsInfo) removeRegionFromSubTreeLocked(region *RegionInfo) {
 	delete(r.subRegions, region.GetMeta().GetId())
 }
 
+// RemoveRegionIfExist removes RegionInfo from regionTree and regionMap if exists.
+func (r *RegionsInfo) RemoveRegionIfExist(id uint64) {
+	if region := r.GetRegion(id); region != nil {
+		r.RemoveRegion(region)
+		r.RemoveRegionFromSubTree(region)
+	}
+}
+
 type peerSlice []*metapb.Peer
 
 func (s peerSlice) Len() int {