From e80022cb8510719d2259a726fd544edb21ed2fce Mon Sep 17 00:00:00 2001
From: Yilin Chen <sticnarf@gmail.com>
Date: Thu, 11 Aug 2022 18:14:55 +0800
Subject: [PATCH] Clear intersecting regions in the cache when inserting a
 region

Signed-off-by: Yilin Chen <sticnarf@gmail.com>
---
 internal/locate/region_cache.go         |  5 +++++
 internal/locate/sorted_btree.go         | 22 ++++++++++++++++++++++
 internal/locate/sorted_btree_generic.go | 21 +++++++++++++++++++++
 3 files changed, 48 insertions(+)

diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index e1558df236..ec2eb5cf95 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -1316,6 +1316,11 @@ func (c *RegionCache) insertRegionToCache(cachedRegion *Region) {
 	if !ok || latest.GetVer() < newVer.GetVer() || latest.GetConfVer() < newVer.GetConfVer() {
 		c.mu.latestVersions[cachedRegion.VerID().id] = newVer
 	}
+	// The intersecting regions in the cache are probably stale, clear them.
+	deleted := c.mu.sorted.removeIntersecting(cachedRegion)
+	for _, r := range deleted {
+		c.removeVersionFromCache(r.cachedRegion.VerID(), r.cachedRegion.GetID())
+	}
 }
 
 // searchCachedRegion finds a region from cache by key. Like `getCachedRegion`,
diff --git a/internal/locate/sorted_btree.go b/internal/locate/sorted_btree.go
index 933332e059..48c4a1d216 100644
--- a/internal/locate/sorted_btree.go
+++ b/internal/locate/sorted_btree.go
@@ -94,6 +94,28 @@ func (s *SortedRegions) AscendGreaterOrEqual(startKey, endKey []byte, limit int)
 	return regions
 }
 
+// removeIntersecting removes all items that have intersection with the key range of given region.
+// If the region itself is in the cache, it's not removed.
+func (s *SortedRegions) removeIntersecting(r *Region) []*btreeItem {
+	var deleted []*btreeItem
+	s.b.AscendGreaterOrEqual(newBtreeSearchItem(r.StartKey()), func(item btree.Item) bool {
+		region := item.(*btreeItem).cachedRegion
+		// Skip the item that is equal to the given region.
+		if region.VerID() == r.VerID() {
+			return true
+		}
+		if len(r.EndKey()) > 0 && bytes.Compare(region.StartKey(), r.EndKey()) >= 0 {
+			return false
+		}
+		deleted = append(deleted, item.(*btreeItem))
+		return true
+	})
+	for _, item := range deleted {
+		s.b.Delete(item)
+	}
+	return deleted
+}
+
 // Clear removes all items from the btree.
 func (s *SortedRegions) Clear() {
 	s.b.Clear(false)
diff --git a/internal/locate/sorted_btree_generic.go b/internal/locate/sorted_btree_generic.go
index 02ca4cd538..4ff61b4361 100644
--- a/internal/locate/sorted_btree_generic.go
+++ b/internal/locate/sorted_btree_generic.go
@@ -94,6 +94,27 @@ func (s *SortedRegions) AscendGreaterOrEqual(startKey, endKey []byte, limit int)
 	return regions
 }
 
+// removeIntersecting removes all items that have intersection with the key range of given region.
+// If the region itself is in the cache, it's not removed.
+func (s *SortedRegions) removeIntersecting(r *Region) []*btreeItem {
+	var deleted []*btreeItem
+	s.b.AscendGreaterOrEqual(newBtreeSearchItem(r.StartKey()), func(item *btreeItem) bool {
+		// Skip the item that is equal to the given region.
+		if item.cachedRegion.VerID() == r.VerID() {
+			return true
+		}
+		if len(r.EndKey()) > 0 && bytes.Compare(item.cachedRegion.StartKey(), r.EndKey()) >= 0 {
+			return false
+		}
+		deleted = append(deleted, item)
+		return true
+	})
+	for _, item := range deleted {
+		s.b.Delete(item)
+	}
+	return deleted
+}
+
 // Clear removes all items from the btree.
 func (s *SortedRegions) Clear() {
 	s.b.Clear(false)