region_cache: fix issue that LocateKey may returns a wrong region (#1299)

* region_cache: fix issue that LocateKey may returns a wrong region

Signed-off-by: zyguan <zhongyangguan@gmail.com>

* address the comment

Signed-off-by: zyguan <zhongyangguan@gmail.com>

---------

Signed-off-by: zyguan <zhongyangguan@gmail.com>

internal/locate/region_cache.go

@@ -1318,18 +1318,10 @@ func (c *RegionCache) findRegionByKey(bo *retry.Backoffer, key []byte, isEndKey
 		}
 	} else if flags := r.resetSyncFlags(needReloadOnAccess | needDelayedReloadReady); flags > 0 {
 		// load region when it be marked as need reload.
-		reloadOnAccess := flags&needReloadOnAccess > 0
+		observeLoadRegion(tag, r, expired, flags)
-		var (
+		// NOTE: we can NOT use c.loadRegionByID(bo, r.GetID()) here because the new region (loaded by id) is not
-			lr  *Region
+		// guaranteed to contain the key. (ref: https://github.com/tikv/client-go/pull/1299)
-			err error
+		lr, err := c.loadRegion(bo, key, isEndKey)
-		)
-		if reloadOnAccess {
-			observeLoadRegion(tag, r, expired, flags)
-			lr, err = c.loadRegion(bo, key, isEndKey)
-		} else {
-			observeLoadRegion("ByID", r, expired, flags)
-			lr, err = c.loadRegionByID(bo, r.GetID())
-		}
 		if err != nil {
 			// ignore error and use old region info.
 			logutil.Logger(bo.GetCtx()).Error("load region failure",
@@ -1337,6 +1329,7 @@ func (c *RegionCache) findRegionByKey(bo *retry.Backoffer, key []byte, isEndKey
 				zap.String("encode-key", util.HexRegionKeyStr(c.codec.EncodeRegionKey(key))))
 		} else {
 			logutil.Eventf(bo.GetCtx(), "load region %d from pd, due to need-reload", lr.GetID())
+			reloadOnAccess := flags&needReloadOnAccess > 0
 			r = lr
 			c.mu.Lock()
 			c.insertRegionToCache(r, reloadOnAccess, reloadOnAccess)

internal/locate/region_cache_test.go

@@ -2204,6 +2204,39 @@ func (s *testRegionCacheSuite) TestRegionCacheHandleHealthStatus() {
 	s.False(store2.healthStatus.IsSlow())
 }
 
+func (s *testRegionCacheSuite) TestSplitThenLocateInvalidRegion() {
+	s.testSplitThenLocateKey(func(r *Region) { r.invalidate(Other) })
+}
+
+func (s *testRegionCacheSuite) TestSplitThenLocateRegionNeedReloadOnAccess() {
+	s.testSplitThenLocateKey(func(r *Region) { r.setSyncFlags(needReloadOnAccess) })
+}
+
+func (s *testRegionCacheSuite) TestSplitThenLocateRegionNeedDelayedReload() {
+	s.testSplitThenLocateKey(func(r *Region) { r.setSyncFlags(needDelayedReloadReady) })
+}
+
+func (s *testRegionCacheSuite) testSplitThenLocateKey(markRegion func(r *Region)) {
+	k := []byte("k")
+
+	// load region to cache
+	_, err := s.cache.LocateRegionByID(s.bo, s.region1)
+	s.NoError(err)
+	r1, expired := s.cache.searchCachedRegionByKey(k, false)
+	s.NotNil(r1)
+	s.False(expired)
+
+	// split region and mark it need sync
+	r2ids := s.cluster.AllocIDs(3)
+	s.cluster.Split(s.region1, r2ids[0], k, r2ids[1:], r2ids[1])
+	markRegion(r1)
+
+	// locate key
+	loc, err := s.cache.LocateKey(s.bo, k)
+	s.NoError(err)
+	s.True(loc.Contains(k))
+}
+
 func (s *testRegionRequestToSingleStoreSuite) TestRefreshCache() {
 	_ = s.cache.refreshRegionIndex(s.bo)
 	r, _ := s.cache.scanRegionsFromCache(s.bo, []byte{}, nil, 10)