Repository: tikv/client-go (go)

Issue:
`insertRegionToCache` should remove intersecting regions in the cache
Our region cache B-Tree uses the start key of the region as the entry key. So, `insertRegionToCache` only replaces the region with the same start key:
https://github.com/tikv/client-go/blob/516cfcdecce865d90ec05dd1cc1de060c35438dc/internal/locate/region_cache.go#L1287

Considering there's a table with lots of regions, the user removes all data from the table. So, PD will finally merge all regions into one. In our region cache implementation, only the first region (containing the table head) is updated, while all other regions still exist in the cache.

When we locate a key near the end of the table, `DescendLessOrEqual` will iterate over all regions that are left in the cache. This will waste a lot of CPU and make region search really slow, and it is also possible to cause lots of unnecessary PD `loadRegion` RPCs.

To solve the problem, when we do `insertRegionToCache`, we should iterate the B-Tree for all regions that intersect with the inserted region and remove them.

This is probably the root cause of https://github.com/tikv/client-go/issues/532

Pull request description:
Clear intersecting regions in the cache when inserting a region
Closes #565

Patch:
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index e1558df..ec2eb5c 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 c72bed6..6165f76 100644
--- a/internal/locate/sorted_btree.go
+++ b/internal/locate/sorted_btree.go
@@ -91,6 +91,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)


Test patch:
diff --git a/internal/locate/region_cache_test.go b/internal/locate/region_cache_test.go
index f995dc4..0914288 100644
--- a/internal/locate/region_cache_test.go
+++ b/internal/locate/region_cache_test.go
@@ -1536,3 +1536,59 @@ func (s *testRegionCacheSuite) TestLocateBucket() {
 		s.True(b.Contains(key))
 	}
 }
+
+func (s *testRegionCacheSuite) TestRemoveIntersectingRegions() {
+	// Split at "b", "c", "d", "e"
+	regions := s.cluster.AllocIDs(4)
+	regions = append([]uint64{s.region1}, regions...)
+
+	peers := [][]uint64{{s.peer1, s.peer2}}
+	for i := 0; i < 4; i++ {
+		peers = append(peers, s.cluster.AllocIDs(2))
+	}
+
+	for i := 0; i < 4; i++ {
+		s.cluster.Split(regions[i], regions[i+1], []byte{'b' + byte(i)}, peers[i+1], peers[i+1][0])
+	}
+
+	for c := 'a'; c <= 'e'; c++ {
+		loc, err := s.cache.LocateKey(s.bo, []byte{byte(c)})
+		s.Nil(err)
+		s.Equal(loc.Region.GetID(), regions[c-'a'])
+	}
+
+	// merge all except the last region together
+	for i := 1; i <= 3; i++ {
+		s.cluster.Merge(regions[0], regions[i])
+	}
+
+	// Now the region cache contains stale information
+	loc, err := s.cache.LocateKey(s.bo, []byte{'c'})
+	s.Nil(err)
+	s.NotEqual(loc.Region.GetID(), regions[0]) // This is incorrect, but is expected
+	loc, err = s.cache.LocateKey(s.bo, []byte{'e'})
+	s.Nil(err)
+	s.Equal(loc.Region.GetID(), regions[4]) // 'e' is not merged yet, so it's still correct
+
+	// If we insert the new region into the cache, the old intersecting regions will be removed.
+	// And the result will be correct.
+	region, err := s.cache.loadRegion(s.bo, []byte("c"), false)
+	s.Nil(err)
+	s.Equal(region.GetID(), regions[0])
+	s.cache.insertRegionToCache(region)
+	loc, err = s.cache.LocateKey(s.bo, []byte{'c'})
+	s.Nil(err)
+	s.Equal(loc.Region.GetID(), regions[0])
+	s.checkCache(2)
+
+	// Now, we merge the last region. This case tests against how we handle the empty end_key.
+	s.cluster.Merge(regions[0], regions[4])
+	region, err = s.cache.loadRegion(s.bo, []byte("e"), false)
+	s.Nil(err)
+	s.Equal(region.GetID(), regions[0])
+	s.cache.insertRegionToCache(region)
+	loc, err = s.cache.LocateKey(s.bo, []byte{'e'})
+	s.Nil(err)
+	s.Equal(loc.Region.GetID(), regions[0])
+	s.checkCache(1)
+}