Skip to content

Commit

Permalink
core: simplify regionSubTree and regionMap (tikv#3693)
Browse files Browse the repository at this point in the history
* core: simplify regionSubTree and regionMap

Signed-off-by: HunDunDM <hundundm@gmail.com>

* fix stat error

Signed-off-by: HunDunDM <hundundm@gmail.com>

Co-authored-by: Ti Chi Robot <ti-community-prow-bot@tidb.io>
  • Loading branch information
HunDunDM and ti-chi-bot authored Jun 2, 2021
1 parent 037d673 commit ab8b5b3
Show file tree
Hide file tree
Showing 4 changed files with 96 additions and 181 deletions.
183 changes: 33 additions & 150 deletions server/core/region.go
Original file line number Diff line number Diff line change
Expand Up @@ -442,165 +442,47 @@ func (r *RegionInfo) GetReplicationStatus() *replication_modepb.RegionReplicatio
}

// regionMap wraps a map[uint64]*core.RegionInfo and supports randomly pick a region.
type regionMap struct {
m map[uint64]*RegionInfo
totalSize int64
totalKeys int64
}
type regionMap map[uint64]*RegionInfo

func newRegionMap() *regionMap {
return &regionMap{
m: make(map[uint64]*RegionInfo),
}
func newRegionMap() regionMap {
return make(map[uint64]*RegionInfo)
}

func (rm *regionMap) Len() int {
if rm == nil {
return 0
}
return len(rm.m)
func (rm regionMap) Len() int {
return len(rm)
}

func (rm *regionMap) Get(id uint64) *RegionInfo {
if rm == nil {
return nil
}
if r, ok := rm.m[id]; ok {
return r
}
return nil
func (rm regionMap) Get(id uint64) *RegionInfo {
return rm[id]
}

func (rm *regionMap) Put(region *RegionInfo) {
if old, ok := rm.m[region.GetID()]; ok {
rm.totalSize -= old.approximateSize
rm.totalKeys -= old.approximateKeys
}
rm.m[region.GetID()] = region
rm.totalSize += region.approximateSize
rm.totalKeys += region.approximateKeys
func (rm regionMap) Put(region *RegionInfo) {
rm[region.GetID()] = region
}

func (rm *regionMap) Delete(id uint64) {
if rm == nil {
return
}
if old, ok := rm.m[id]; ok {
delete(rm.m, id)
rm.totalSize -= old.approximateSize
rm.totalKeys -= old.approximateKeys
}
}

func (rm *regionMap) TotalSize() int64 {
if rm.Len() == 0 {
return 0
}
return rm.totalSize
}

// regionSubTree is used to manager different types of regions.
type regionSubTree struct {
*regionTree
totalSize int64
totalKeys int64
}

func newRegionSubTree() *regionSubTree {
return &regionSubTree{
regionTree: newRegionTree(),
totalSize: 0,
}
}

func (rst *regionSubTree) TotalSize() int64 {
if rst.length() == 0 {
return 0
}
return rst.totalSize
}

func (rst *regionSubTree) scanRanges() []*RegionInfo {
if rst.length() == 0 {
return nil
}
var res []*RegionInfo
rst.scanRange([]byte(""), func(region *RegionInfo) bool {
res = append(res, region)
return true
})
return res
}

func (rst *regionSubTree) update(region *RegionInfo) {
overlaps := rst.regionTree.update(region)
rst.totalSize += region.approximateSize
rst.totalKeys += region.approximateKeys
for _, r := range overlaps {
rst.totalSize -= r.approximateSize
rst.totalKeys -= r.approximateKeys
}
}

func (rst *regionSubTree) remove(region *RegionInfo) {
if rst.length() == 0 {
return
}
if rst.regionTree.remove(region) != nil {
rst.totalSize -= region.approximateSize
rst.totalKeys -= region.approximateKeys
}
}

func (rst *regionSubTree) length() int {
if rst == nil {
return 0
}
return rst.regionTree.length()
}

func (rst *regionSubTree) RandomRegion(ranges []KeyRange) *RegionInfo {
if rst.length() == 0 {
return nil
}

return rst.regionTree.RandomRegion(ranges)
}

func (rst *regionSubTree) RandomRegions(n int, ranges []KeyRange) []*RegionInfo {
if rst.length() == 0 {
return nil
}

regions := make([]*RegionInfo, 0, n)

for i := 0; i < n; i++ {
if region := rst.regionTree.RandomRegion(ranges); region != nil {
regions = append(regions, region)
}
}
return regions
func (rm regionMap) Delete(id uint64) {
delete(rm, id)
}

// RegionsInfo for export
type RegionsInfo struct {
tree *regionTree
regions *regionMap // regionID -> regionInfo
leaders map[uint64]*regionSubTree // storeID -> regionSubTree
followers map[uint64]*regionSubTree // storeID -> regionSubTree
learners map[uint64]*regionSubTree // storeID -> regionSubTree
pendingPeers map[uint64]*regionSubTree // storeID -> regionSubTree
regions regionMap // regionID -> regionInfo
leaders map[uint64]*regionTree // storeID -> sub regionTree
followers map[uint64]*regionTree // storeID -> sub regionTree
learners map[uint64]*regionTree // storeID -> sub regionTree
pendingPeers map[uint64]*regionTree // storeID -> sub regionTree
}

// NewRegionsInfo creates RegionsInfo with tree, regions, leaders and followers
func NewRegionsInfo() *RegionsInfo {
return &RegionsInfo{
tree: newRegionTree(),
regions: newRegionMap(),
leaders: make(map[uint64]*regionSubTree),
followers: make(map[uint64]*regionSubTree),
learners: make(map[uint64]*regionSubTree),
pendingPeers: make(map[uint64]*regionSubTree),
leaders: make(map[uint64]*regionTree),
followers: make(map[uint64]*regionTree),
learners: make(map[uint64]*regionTree),
pendingPeers: make(map[uint64]*regionTree),
}
}

Expand All @@ -626,13 +508,13 @@ func (r *RegionsInfo) SetRegion(region *RegionInfo) []*RegionInfo {
return r.AddRegion(region)
}

// Length returns the RegionsInfo length
func (r *RegionsInfo) Length() int {
// Len returns the RegionsInfo length
func (r *RegionsInfo) Len() int {
return r.regions.Len()
}

// TreeLength returns the RegionsInfo tree length(now only used in test)
func (r *RegionsInfo) TreeLength() int {
// TreeLen returns the RegionsInfo tree length(now only used in test)
func (r *RegionsInfo) TreeLen() int {
return r.tree.length()
}

Expand All @@ -653,6 +535,7 @@ func (r *RegionsInfo) AddRegion(region *RegionInfo) []*RegionInfo {
if bytes.Equal(regionOld.region.GetStartKey(), region.GetStartKey()) &&
bytes.Equal(regionOld.region.GetEndKey(), region.GetEndKey()) &&
regionOld.region.GetID() == region.GetID() {
r.tree.updateStat(regionOld.region, region)
regionOld.region = region
treeNeedAdd = false
}
Expand All @@ -675,15 +558,15 @@ func (r *RegionsInfo) AddRegion(region *RegionInfo) []*RegionInfo {
// Add leader peer to leaders.
store, ok := r.leaders[storeID]
if !ok {
store = newRegionSubTree()
store = newRegionTree()
r.leaders[storeID] = store
}
store.update(region)
} else {
// Add follower peer to followers.
store, ok := r.followers[storeID]
if !ok {
store = newRegionSubTree()
store = newRegionTree()
r.followers[storeID] = store
}
store.update(region)
Expand All @@ -695,7 +578,7 @@ func (r *RegionsInfo) AddRegion(region *RegionInfo) []*RegionInfo {
storeID := peer.GetStoreId()
store, ok := r.learners[storeID]
if !ok {
store = newRegionSubTree()
store = newRegionTree()
r.learners[storeID] = store
}
store.update(region)
Expand All @@ -705,7 +588,7 @@ func (r *RegionsInfo) AddRegion(region *RegionInfo) []*RegionInfo {
storeID := peer.GetStoreId()
store, ok := r.pendingPeers[storeID]
if !ok {
store = newRegionSubTree()
store = newRegionTree()
r.pendingPeers[storeID] = store
}
store.update(region)
Expand Down Expand Up @@ -824,7 +707,7 @@ func (r *RegionsInfo) SearchPrevRegion(regionKey []byte) *RegionInfo {
// GetRegions gets all RegionInfo from regionMap
func (r *RegionsInfo) GetRegions() []*RegionInfo {
regions := make([]*RegionInfo, 0, r.regions.Len())
for _, region := range r.regions.m {
for _, region := range r.regions {
regions = append(regions, region)
}
return regions
Expand Down Expand Up @@ -868,7 +751,7 @@ func (r *RegionsInfo) GetStoreRegionSize(storeID uint64) int64 {
// GetMetaRegions gets a set of metapb.Region from regionMap
func (r *RegionsInfo) GetMetaRegions() []*metapb.Region {
regions := make([]*metapb.Region, 0, r.regions.Len())
for _, region := range r.regions.m {
for _, region := range r.regions {
regions = append(regions, proto.Clone(region.meta).(*metapb.Region))
}
return regions
Expand Down Expand Up @@ -1019,10 +902,10 @@ func (r *RegionsInfo) GetAdjacentRegions(region *RegionInfo) (*RegionInfo, *Regi

// GetAverageRegionSize returns the average region approximate size.
func (r *RegionsInfo) GetAverageRegionSize() int64 {
if r.regions.Len() == 0 {
if r.tree.length() == 0 {
return 0
}
return r.regions.TotalSize() / int64(r.regions.Len())
return r.tree.TotalSize() / int64(r.tree.length())
}

// DiffRegionPeersInfo return the difference of peers info between two RegionInfo
Expand Down
17 changes: 3 additions & 14 deletions server/core/region_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -160,10 +160,6 @@ var _ = Suite(&testRegionMapSuite{})
type testRegionMapSuite struct{}

func (s *testRegionMapSuite) TestRegionMap(c *C) {
var empty *regionMap
c.Assert(empty.Len(), Equals, 0)
c.Assert(empty.Get(1), IsNil)

rm := newRegionMap()
s.check(c, rm)
rm.Put(s.regionInfo(1))
Expand Down Expand Up @@ -195,7 +191,7 @@ func (s *testRegionMapSuite) regionInfo(id uint64) *RegionInfo {
}
}

func (s *testRegionMapSuite) check(c *C, rm *regionMap, ids ...uint64) {
func (s *testRegionMapSuite) check(c *C, rm regionMap, ids ...uint64) {
// Check Get.
for _, id := range ids {
c.Assert(rm.Get(id).GetID(), Equals, id)
Expand All @@ -208,16 +204,10 @@ func (s *testRegionMapSuite) check(c *C, rm *regionMap, ids ...uint64) {
expect[id] = struct{}{}
}
set1 := make(map[uint64]struct{})
for _, r := range rm.m {
for _, r := range rm {
set1[r.GetID()] = struct{}{}
}
c.Assert(set1, DeepEquals, expect)
// Check region size.
var total int64
for _, id := range ids {
total += int64(id)
}
c.Assert(rm.TotalSize(), Equals, total)
}

var _ = Suite(&testRegionKey{})
Expand Down Expand Up @@ -325,8 +315,7 @@ func (*testRegionKey) TestSetRegion(c *C) {
c.Assert(regions.tree.length(), Equals, 96)
c.Assert(len(regions.GetRegions()), Equals, 96)
c.Assert(regions.GetRegion(201), NotNil)
c.Assert(regions.regions.totalKeys, Equals, int64(20))
c.Assert(regions.regions.totalSize, Equals, int64(30))
c.Assert(regions.tree.TotalSize(), Equals, int64(30))
}

func (*testRegionKey) TestShouldRemoveFromSubTree(c *C) {
Expand Down
Loading

0 comments on commit ab8b5b3

Please sign in to comment.