Wrap locateBucket to handle with edges (#486)

Signed-off-by: SpadeA-Tang <u6748471@anu.edu.au>

internal/locate/region_cache.go

@@ -755,6 +755,40 @@ func (l *KeyLocation) GetBucketVersion() uint64 {
 	return l.Buckets.GetVersion()
 }
 
+// LocateBucketV2 will not return nil if the key is in the region.
+// LocateBucketV2 is similar with LocateBucket. The difference is that when the key is in [KeyLocation.StartKey, first Bucket key)
+// it will return Bucket{KeyLocation.StartKey, first Bucket key} rather than nil --- it's reasonable to assume that
+// Bucket{KeyLocation.StartKey, first Bucket key} is a bucket belonging to the region. Key in [last Bucket key, KeyLocation.EndKey)
+// is handled similarly.
+func (l *KeyLocation) LocateBucketV2(key []byte) *Bucket {
+	bucket := l.LocateBucket(key)
+	if bucket != nil || !l.Contains(key) {
+		return bucket
+	}
+	counts := len(l.Buckets.Keys)
+	if counts == 0 {
+		return &Bucket{
+			l.StartKey,
+			l.EndKey,
+		}
+	}
+	firstBucketKey := l.Buckets.Keys[0]
+	if bytes.Compare(key, firstBucketKey) < 0 {
+		return &Bucket{
+			l.StartKey,
+			firstBucketKey,
+		}
+	}
+	lastBucketKey := l.Buckets.Keys[counts-1]
+	if bytes.Compare(lastBucketKey, key) <= 0 {
+		return &Bucket{
+			lastBucketKey,
+			l.EndKey,
+		}
+	}
+	return bucket
+}
+
 // LocateBucket returns the bucket the key is located.
 func (l *KeyLocation) LocateBucket(key []byte) *Bucket {
 	keys := l.Buckets.GetKeys()

internal/locate/region_cache_test.go

@@ -1505,3 +1505,47 @@ func (s *testRegionCacheSuite) TestBuckets() {
 	s.cache.UpdateBucketsIfNeeded(cachedRegion.VerID(), newBuckets.GetVersion())
 	waitUpdateBuckets(newBuckets, []byte("a"))
 }
+
+func (s *testRegionCacheSuite) TestLocateBucketV2() {
+	// proto.Clone clones []byte{} to nil and [][]byte{nil or []byte{}} to [][]byte{[]byte{}}.
+	// nilToEmtpyBytes unifies it for tests.
+	nilToEmtpyBytes := func(s []byte) []byte {
+		if s == nil {
+			s = []byte{}
+		}
+		return s
+	}
+	r, _ := s.cluster.GetRegion(s.region1)
+
+	// First test normal case: region start equals to the first bucket keys and
+	// region end equals to the last bucket key
+	bucketKeys := [][]byte{nilToEmtpyBytes(r.GetStartKey()), []byte("a"), []byte("b"), nilToEmtpyBytes(r.GetEndKey())}
+	s.cluster.SplitRegionBuckets(s.region1, bucketKeys, uint64(time.Now().Nanosecond()))
+	loc, err := s.cache.LocateKey(s.bo, []byte("a"))
+	s.NotNil(loc)
+	s.Nil(err)
+	for _, key := range [][]byte{{}, {'a' - 1}, []byte("a"), []byte("a0"), []byte("b"), []byte("c")} {
+		b := loc.LocateBucket(key)
+		s.NotNil(b)
+		s.True(b.Contains(key))
+	}
+
+	// Then test cases where there's some holes in region start and the first bucket key
+	// and in the last bucket key and region end
+	bucketKeys = [][]byte{[]byte("a"), []byte("b")}
+	bucketVersion := uint64(time.Now().Nanosecond())
+	s.cluster.SplitRegionBuckets(s.region1, bucketKeys, bucketVersion)
+	s.cache.UpdateBucketsIfNeeded(s.getRegion([]byte("a")).VerID(), bucketVersion)
+	// wait for region update
+	time.Sleep(300 * time.Millisecond)
+	loc, err = s.cache.LocateKey(s.bo, []byte("a"))
+	s.NotNil(loc)
+	s.Nil(err)
+	for _, key := range [][]byte{{'a' - 1}, []byte("c")} {
+		b := loc.LocateBucket(key)
+		s.Nil(b)
+		b = loc.LocateBucketV2(key)
+		s.NotNil(b)
+		s.True(b.Contains(key))
+	}
+}