address the comments

Signed-off-by: Ryan Leung <rleungx@gmail.com>
tikv · Oct 11, 2019 · 46a47f6 · 46a47f6
1 parent bda04dc
commit 46a47f6
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 51 deletions.
diff --git a/server/core/region.go b/server/core/region.go
@@ -17,7 +17,6 @@ import (
 	"bytes"
 	"encoding/hex"
 	"fmt"
-	"math/rand"
 	"reflect"
 	"strings"
 
@@ -520,37 +519,7 @@ func (rst *regionSubTree) RandomRegion(startKey, endKey []byte) *RegionInfo {
 	}
 
 	t := rst.regionTree
-	var startIndex, endIndex, index int
-	var startRegion, endRegion *RegionInfo
-
-	if len(startKey) != 0 {
-		startRegion, startIndex = t.getWithIndex(&regionItem{region: &RegionInfo{meta: &metapb.Region{StartKey: startKey}}})
-	} else {
-		startRegion, startIndex = t.getWithIndex(t.tree.Min())
-	}
-
-	if len(endKey) != 0 {
-		endRegion, endIndex = t.getWithIndex(&regionItem{region: &RegionInfo{meta: &metapb.Region{StartKey: endKey}}})
-	} else {
-		_, endIndex = t.getWithIndex(t.tree.Max())
-		endRegion = nil
-		endIndex++
-	}
-
-	if endIndex == startIndex {
-		if endRegion == nil {
-			return t.tree.GetAt(startIndex - 1).(*regionItem).region
-		}
-		return t.tree.GetAt(startIndex).(*regionItem).region
-	}
-
-	if endRegion == nil && startRegion == nil {
-		index = rand.Intn(endIndex-startIndex+1) + startIndex
-		return t.tree.GetAt(index - 1).(*regionItem).region
-	}
-
-	index = rand.Intn(endIndex-startIndex) + startIndex
-	return t.tree.GetAt(index).(*regionItem).region
+	return t.RandomRegion(startKey, endKey)
 }
 
 // RegionsInfo for export

diff --git a/server/core/region_tree.go b/server/core/region_tree.go
@@ -194,6 +194,7 @@ func (t *regionTree) getAdjacentRegions(region *RegionInfo) (*regionItem, *regio
 	return prev, next
 }
 
+// RandomRegion is used to get a random region from a given range [startKey, endKey].
 func (t *regionTree) RandomRegion(startKey, endKey []byte) *RegionInfo {
 	if t.length() == 0 {
 		return nil
@@ -204,36 +205,27 @@ func (t *regionTree) RandomRegion(startKey, endKey []byte) *RegionInfo {
 		return t.tree.Min().(*regionItem).region
 	}
 
-	var startIndex, endIndex, index int
-	var startRegion, endRegion *RegionInfo
-
+	var (
+		startRegion          *RegionInfo
+		startIndex, endIndex int
+	)
 	if len(startKey) != 0 {
 		startRegion, startIndex = t.getWithIndex(&regionItem{region: &RegionInfo{meta: &metapb.Region{StartKey: startKey}}})
 	} else {
-		startRegion, startIndex = t.getWithIndex(t.tree.Min())
+		startRegion, startIndex = t.getWithIndex(t.tree.Min()) // for test purpose
 	}
 
 	if len(endKey) != 0 {
-		endRegion, endIndex = t.getWithIndex(&regionItem{region: &RegionInfo{meta: &metapb.Region{StartKey: endKey}}})
+		_, endIndex = t.getWithIndex(&regionItem{region: &RegionInfo{meta: &metapb.Region{StartKey: endKey}}})
 	} else {
-		_, endIndex = t.getWithIndex(t.tree.Max())
-		endRegion = nil
-		endIndex++
-	}
-
-	if endIndex == startIndex {
-		if endRegion == nil {
-			return t.tree.GetAt(startIndex - 1).(*regionItem).region
-		}
-		return t.tree.GetAt(startIndex).(*regionItem).region
+		endIndex = t.tree.Len()
 	}
 
-	if endRegion == nil && startRegion == nil {
-		index = rand.Intn(endIndex-startIndex+1) + startIndex
-		return t.tree.GetAt(index - 1).(*regionItem).region
+	if startRegion == nil {
+		startIndex--
 	}
 
-	index = rand.Intn(endIndex-startIndex) + startIndex
+	index := rand.Intn(endIndex-startIndex) + startIndex
 	return t.tree.GetAt(index).(*regionItem).region
 }
 

diff --git a/server/core/region_tree_test.go b/server/core/region_tree_test.go
@@ -236,6 +236,8 @@ func (s *testRegionSuite) TestRegionTreeSplitAndMerge(c *C) {
 
 func (s *testRegionSuite) TestRandomRegion(c *C) {
 	tree := newRegionTree()
+	r := tree.RandomRegion([]byte(""), []byte(""))
+	c.Assert(r, IsNil)
 
 	regionA := NewTestRegionInfo([]byte(""), []byte("g"))
 	tree.update(regionA)