Skip to content

Commit

Permalink
mockcluster supports tiflash learner
Browse files Browse the repository at this point in the history 
Signed-off-by: HunDunDM <hundundm@gmail.com>
  • Loading branch information
@HunDunDM
HunDunDM committed Aug 6, 2021
1 parent d776340 commit 89a9d1d
Showing 1 changed file with 29 additions and 20 deletions.
49 changes: 29 additions & 20 deletions pkg/mock/mockcluster/mockcluster.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,9 +85,9 @@ func (mc *Cluster) ScanRegions(startKey, endKey []byte, limit int) []*core.Regio
}

// LoadRegion puts region info without leader
func (mc *Cluster) LoadRegion(regionID uint64, followerIds ...uint64) {
func (mc *Cluster) LoadRegion(regionID uint64, peerStoreIDs ...uint64) {
// regions load from etcd will have no leader
r := mc.newMockRegionInfo(regionID, 0, followerIds...).Clone(core.WithLeader(nil))
r := mc.newMockRegionInfo(regionID, 0, peerStoreIDs...).Clone(core.WithLeader(nil))
mc.PutRegion(r)
}

Expand Down Expand Up @@ -304,8 +304,8 @@ func (mc *Cluster) AddLabelsStore(storeID uint64, regionCount int, labels map[st
}

// AddLeaderRegion adds region with specified leader and followers.
func (mc *Cluster) AddLeaderRegion(regionID uint64, leaderStoreID uint64, followerStoreIDs ...uint64) *core.RegionInfo {
origin := mc.newMockRegionInfo(regionID, leaderStoreID, followerStoreIDs...)
func (mc *Cluster) AddLeaderRegion(regionID uint64, leaderStoreID uint64, otherPeerStoreIDs ...uint64) *core.RegionInfo {
origin := mc.newMockRegionInfo(regionID, leaderStoreID, otherPeerStoreIDs...)
region := origin.Clone(core.SetApproximateSize(defaultRegionSize/mb), core.SetApproximateKeys(10))
mc.PutRegion(region)
return region
Expand All @@ -320,8 +320,8 @@ func (mc *Cluster) AddRegionWithLearner(regionID uint64, leaderStoreID uint64, f
}

// AddLeaderRegionWithRange adds region with specified leader, followers and key range.
func (mc *Cluster) AddLeaderRegionWithRange(regionID uint64, startKey string, endKey string, leaderID uint64, followerIds ...uint64) {
o := mc.newMockRegionInfo(regionID, leaderID, followerIds...)
func (mc *Cluster) AddLeaderRegionWithRange(regionID uint64, startKey string, endKey string, LeaderStoreID uint64, otherPeerStoreIDs ...uint64) {
o := mc.newMockRegionInfo(regionID, LeaderStoreID, otherPeerStoreIDs...)
r := o.Clone(
core.WithStartKey([]byte(startKey)),
core.WithEndKey([]byte(endKey)),
Expand All @@ -331,11 +331,11 @@ func (mc *Cluster) AddLeaderRegionWithRange(regionID uint64, startKey string, en

// AddRegionWithReadInfo adds region with specified leader, followers and read info.
func (mc *Cluster) AddRegionWithReadInfo(
regionID uint64, leaderID uint64,
regionID uint64, LeaderStoreID uint64,
readBytes, readKeys uint64,
reportInterval uint64,
followerIds []uint64, filledNums ...int) []*statistics.HotPeerStat {
r := mc.newMockRegionInfo(regionID, leaderID, followerIds...)
otherPeerStoreIDs []uint64, filledNums ...int) []*statistics.HotPeerStat {
r := mc.newMockRegionInfo(regionID, LeaderStoreID, otherPeerStoreIDs...)
r = r.Clone(core.SetReadBytes(readBytes))
r = r.Clone(core.SetReadKeys(readKeys))
r = r.Clone(core.SetReportInterval(reportInterval))
Expand All @@ -356,9 +356,9 @@ func (mc *Cluster) AddRegionWithReadInfo(
}

// AddRegionWithPeerReadInfo adds region with specified peer read info.
func (mc *Cluster) AddRegionWithPeerReadInfo(regionID, leaderID, targetStoreID, readBytes, readKeys, reportInterval uint64,
followerIds []uint64, filledNums ...int) []*statistics.HotPeerStat {
r := mc.newMockRegionInfo(regionID, leaderID, followerIds...)
func (mc *Cluster) AddRegionWithPeerReadInfo(regionID, LeaderStoreID, targetStoreID, readBytes, readKeys, reportInterval uint64,
otherPeerStoreIDs []uint64, filledNums ...int) []*statistics.HotPeerStat {
r := mc.newMockRegionInfo(regionID, LeaderStoreID, otherPeerStoreIDs...)
r = r.Clone(core.SetReadBytes(readBytes), core.SetReadKeys(readKeys), core.SetReportInterval(reportInterval))
filledNum := mc.HotCache.GetFilledPeriod(statistics.ReadFlow)
if len(filledNums) > 0 {
Expand All @@ -379,11 +379,11 @@ func (mc *Cluster) AddRegionWithPeerReadInfo(regionID, leaderID, targetStoreID,

// AddRegionLeaderWithReadInfo add region leader read info
func (mc *Cluster) AddRegionLeaderWithReadInfo(
regionID uint64, leaderID uint64,
regionID uint64, LeaderStoreID uint64,
readBytes, readKeys uint64,
reportInterval uint64,
followerIds []uint64, filledNums ...int) []*statistics.HotPeerStat {
r := mc.newMockRegionInfo(regionID, leaderID, followerIds...)
otherPeerStoreIDs []uint64, filledNums ...int) []*statistics.HotPeerStat {
r := mc.newMockRegionInfo(regionID, LeaderStoreID, otherPeerStoreIDs...)
r = r.Clone(core.SetReadBytes(readBytes))
r = r.Clone(core.SetReadKeys(readKeys))
r = r.Clone(core.SetReportInterval(reportInterval))
Expand All @@ -405,11 +405,11 @@ func (mc *Cluster) AddRegionLeaderWithReadInfo(

// AddLeaderRegionWithWriteInfo adds region with specified leader and peers write info.
func (mc *Cluster) AddLeaderRegionWithWriteInfo(
regionID uint64, leaderID uint64,
regionID uint64, LeaderStoreID uint64,
writtenBytes, writtenKeys uint64,
reportInterval uint64,
followerIds []uint64, filledNums ...int) []*statistics.HotPeerStat {
r := mc.newMockRegionInfo(regionID, leaderID, followerIds...)
otherPeerStoreIDs []uint64, filledNums ...int) []*statistics.HotPeerStat {
r := mc.newMockRegionInfo(regionID, LeaderStoreID, otherPeerStoreIDs...)
r = r.Clone(core.SetWrittenBytes(writtenBytes))
r = r.Clone(core.SetWrittenKeys(writtenKeys))
r = r.Clone(core.SetReportInterval(reportInterval))
Expand Down Expand Up @@ -623,8 +623,17 @@ func (mc *Cluster) UpdateStoreStatus(id uint64) {
mc.PutStore(newStore)
}

func (mc *Cluster) newMockRegionInfo(regionID uint64, leaderStoreID uint64, followerStoreIDs ...uint64) *core.RegionInfo {
return mc.MockRegionInfo(regionID, leaderStoreID, followerStoreIDs, []uint64{}, nil)
func (mc *Cluster) newMockRegionInfo(regionID uint64, leaderStoreID uint64, otherPeerStoreIDs ...uint64) *core.RegionInfo {
var followerStoreIDs []uint64
var learnerStoreIDs []uint64
for _, storeID := range otherPeerStoreIDs {
if store := mc.GetStore(storeID); store != nil && core.IsTiFlashStore(store.GetMeta()) {
learnerStoreIDs = append(learnerStoreIDs, storeID)
} else {
followerStoreIDs = append(followerStoreIDs, storeID)
}
}
return mc.MockRegionInfo(regionID, leaderStoreID, followerStoreIDs, learnerStoreIDs, nil)
}

// CheckLabelProperty checks label property.
Expand Down

0 comments on commit 89a9d1d

Please sign in to comment.