talent-plan-tinykv/scheduler/server/core/region.go
2022-01-11 11:31:44 +08:00

824 lines
22 KiB
Go

// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.
package core
import (
"bytes"
"encoding/hex"
"fmt"
"reflect"
"strings"
"github.com/gogo/protobuf/proto"
"github.com/pingcap-incubator/tinykv/proto/pkg/metapb"
"github.com/pingcap-incubator/tinykv/proto/pkg/schedulerpb"
)
// RegionInfo records detail region info.
// Read-Only once created.
type RegionInfo struct {
meta *metapb.Region
learners []*metapb.Peer
voters []*metapb.Peer
leader *metapb.Peer
pendingPeers []*metapb.Peer
approximateSize int64
}
// NewRegionInfo creates RegionInfo with region's meta and leader peer.
func NewRegionInfo(region *metapb.Region, leader *metapb.Peer, opts ...RegionCreateOption) *RegionInfo {
regionInfo := &RegionInfo{
meta: region,
leader: leader,
}
for _, opt := range opts {
opt(regionInfo)
}
classifyVoterAndLearner(regionInfo)
return regionInfo
}
// classifyVoterAndLearner sorts out voter and learner from peers into different slice.
func classifyVoterAndLearner(region *RegionInfo) {
voters := make([]*metapb.Peer, 0, len(region.meta.Peers))
for _, p := range region.meta.Peers {
voters = append(voters, p)
}
region.voters = voters
}
// EmptyRegionApproximateSize is the region approximate size of an empty region
// (heartbeat size <= 1MB).
const EmptyRegionApproximateSize = 1
// RegionFromHeartbeat constructs a Region from region heartbeat.
func RegionFromHeartbeat(heartbeat *schedulerpb.RegionHeartbeatRequest) *RegionInfo {
// Convert unit to MB.
// If region is empty or less than 1MB, use 1MB instead.
regionSize := heartbeat.GetApproximateSize() / (1 << 20)
if regionSize < EmptyRegionApproximateSize {
regionSize = EmptyRegionApproximateSize
}
region := &RegionInfo{
meta: heartbeat.GetRegion(),
leader: heartbeat.GetLeader(),
pendingPeers: heartbeat.GetPendingPeers(),
approximateSize: int64(regionSize),
}
classifyVoterAndLearner(region)
return region
}
// Clone returns a copy of current regionInfo.
func (r *RegionInfo) Clone(opts ...RegionCreateOption) *RegionInfo {
pendingPeers := make([]*metapb.Peer, 0, len(r.pendingPeers))
for _, peer := range r.pendingPeers {
pendingPeers = append(pendingPeers, proto.Clone(peer).(*metapb.Peer))
}
region := &RegionInfo{
meta: proto.Clone(r.meta).(*metapb.Region),
leader: proto.Clone(r.leader).(*metapb.Peer),
pendingPeers: pendingPeers,
approximateSize: r.approximateSize,
}
for _, opt := range opts {
opt(region)
}
classifyVoterAndLearner(region)
return region
}
// GetLearners returns the learners.
func (r *RegionInfo) GetLearners() []*metapb.Peer {
return r.learners
}
// GetVoters returns the voters.
func (r *RegionInfo) GetVoters() []*metapb.Peer {
return r.voters
}
// GetPeer returns the peer with specified peer id.
func (r *RegionInfo) GetPeer(peerID uint64) *metapb.Peer {
for _, peer := range r.meta.GetPeers() {
if peer.GetId() == peerID {
return peer
}
}
return nil
}
// GetDownLearner returns the down learner with soecified peer id.
func (r *RegionInfo) GetDownLearner(peerID uint64) *metapb.Peer {
return nil
}
// GetPendingPeer returns the pending peer with specified peer id.
func (r *RegionInfo) GetPendingPeer(peerID uint64) *metapb.Peer {
for _, peer := range r.pendingPeers {
if peer.GetId() == peerID {
return peer
}
}
return nil
}
// GetPendingVoter returns the pending voter with specified peer id.
func (r *RegionInfo) GetPendingVoter(peerID uint64) *metapb.Peer {
for _, peer := range r.pendingPeers {
if peer.GetId() == peerID {
return peer
}
}
return nil
}
// GetPendingLearner returns the pending learner peer with specified peer id.
func (r *RegionInfo) GetPendingLearner(peerID uint64) *metapb.Peer {
return nil
}
// GetStorePeer returns the peer in specified store.
func (r *RegionInfo) GetStorePeer(storeID uint64) *metapb.Peer {
for _, peer := range r.meta.GetPeers() {
if peer.GetStoreId() == storeID {
return peer
}
}
return nil
}
// GetStoreVoter returns the voter in specified store.
func (r *RegionInfo) GetStoreVoter(storeID uint64) *metapb.Peer {
for _, peer := range r.voters {
if peer.GetStoreId() == storeID {
return peer
}
}
return nil
}
// GetStoreLearner returns the learner peer in specified store.
func (r *RegionInfo) GetStoreLearner(storeID uint64) *metapb.Peer {
for _, peer := range r.learners {
if peer.GetStoreId() == storeID {
return peer
}
}
return nil
}
// GetStoreIds returns a map indicate the region distributed.
func (r *RegionInfo) GetStoreIds() map[uint64]struct{} {
peers := r.meta.GetPeers()
stores := make(map[uint64]struct{}, len(peers))
for _, peer := range peers {
stores[peer.GetStoreId()] = struct{}{}
}
return stores
}
// GetFollowers returns a map indicate the follow peers distributed.
func (r *RegionInfo) GetFollowers() map[uint64]*metapb.Peer {
peers := r.GetVoters()
followers := make(map[uint64]*metapb.Peer, len(peers))
for _, peer := range peers {
if r.leader == nil || r.leader.GetId() != peer.GetId() {
followers[peer.GetStoreId()] = peer
}
}
return followers
}
// GetFollower randomly returns a follow peer.
func (r *RegionInfo) GetFollower() *metapb.Peer {
for _, peer := range r.GetVoters() {
if r.leader == nil || r.leader.GetId() != peer.GetId() {
return peer
}
}
return nil
}
// GetDiffFollowers returns the followers which is not located in the same
// store as any other followers of the another specified region.
func (r *RegionInfo) GetDiffFollowers(other *RegionInfo) []*metapb.Peer {
res := make([]*metapb.Peer, 0, len(r.meta.Peers))
for _, p := range r.GetFollowers() {
diff := true
for _, o := range other.GetFollowers() {
if p.GetStoreId() == o.GetStoreId() {
diff = false
break
}
}
if diff {
res = append(res, p)
}
}
return res
}
// GetID returns the ID of the region.
func (r *RegionInfo) GetID() uint64 {
return r.meta.GetId()
}
// GetMeta returns the meta information of the region.
func (r *RegionInfo) GetMeta() *metapb.Region {
return r.meta
}
// GetApproximateSize returns the approximate size of the region.
func (r *RegionInfo) GetApproximateSize() int64 {
return r.approximateSize
}
// GetPendingPeers returns the pending peers of the region.
func (r *RegionInfo) GetPendingPeers() []*metapb.Peer {
return r.pendingPeers
}
// GetLeader returns the leader of the region.
func (r *RegionInfo) GetLeader() *metapb.Peer {
return r.leader
}
// GetStartKey returns the start key of the region.
func (r *RegionInfo) GetStartKey() []byte {
return r.meta.StartKey
}
// GetEndKey returns the end key of the region.
func (r *RegionInfo) GetEndKey() []byte {
return r.meta.EndKey
}
// GetPeers returns the peers of the region.
func (r *RegionInfo) GetPeers() []*metapb.Peer {
return r.meta.GetPeers()
}
// GetRegionEpoch returns the region epoch of the region.
func (r *RegionInfo) GetRegionEpoch() *metapb.RegionEpoch {
return r.meta.RegionEpoch
}
// regionMap wraps a map[uint64]*core.RegionInfo and supports randomly pick a region.
type regionMap struct {
m map[uint64]*RegionInfo
totalSize int64
totalKeys int64
}
func newRegionMap() *regionMap {
return &regionMap{
m: make(map[uint64]*RegionInfo),
}
}
func (rm *regionMap) Len() int {
if rm == nil {
return 0
}
return len(rm.m)
}
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) Put(region *RegionInfo) {
if old, ok := rm.m[region.GetID()]; ok {
rm.totalSize -= old.approximateSize
}
rm.m[region.GetID()] = region
rm.totalSize += region.approximateSize
}
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
}
}
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
}
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) {
if r := rst.find(region); r != nil {
rst.totalSize += region.approximateSize - r.region.approximateSize
r.region = region
return
}
rst.totalSize += region.approximateSize
rst.regionTree.update(region)
}
func (rst *regionSubTree) remove(region *RegionInfo) {
if rst.length() == 0 {
return
}
rst.regionTree.remove(region)
}
func (rst *regionSubTree) length() int {
if rst == nil {
return 0
}
return rst.regionTree.length()
}
func (rst *regionSubTree) RandomRegion(startKey, endKey []byte) *RegionInfo {
if rst.length() == 0 {
return nil
}
return rst.regionTree.RandomRegion(startKey, endKey)
}
// 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
}
// 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),
}
}
// GetRegion returns the RegionInfo with regionID
func (r *RegionsInfo) GetRegion(regionID uint64) *RegionInfo {
region := r.regions.Get(regionID)
if region == nil {
return nil
}
return region
}
// SetRegion sets the RegionInfo with regionID
func (r *RegionsInfo) SetRegion(region *RegionInfo) []*RegionInfo {
if origin := r.regions.Get(region.GetID()); origin != nil {
r.RemoveRegion(origin)
}
return r.AddRegion(region)
}
// Length returns the RegionsInfo length
func (r *RegionsInfo) Length() int {
return r.regions.Len()
}
// TreeLength returns the RegionsInfo tree length(now only used in test)
func (r *RegionsInfo) TreeLength() int {
return r.tree.length()
}
// GetOverlaps returns the regions which are overlapped with the specified region range.
func (r *RegionsInfo) GetOverlaps(region *RegionInfo) []*RegionInfo {
return r.tree.getOverlaps(region)
}
// AddRegion adds RegionInfo to regionTree and regionMap, also update leaders and followers by region peers
func (r *RegionsInfo) AddRegion(region *RegionInfo) []*RegionInfo {
// Add to tree and regions.
overlaps := r.tree.update(region)
for _, item := range overlaps {
r.RemoveRegion(r.GetRegion(item.GetID()))
}
r.regions.Put(region)
// Add to leaders and followers.
for _, peer := range region.GetVoters() {
storeID := peer.GetStoreId()
if peer.GetId() == region.leader.GetId() {
// Add leader peer to leaders.
store, ok := r.leaders[storeID]
if !ok {
store = newRegionSubTree()
r.leaders[storeID] = store
}
store.update(region)
} else {
// Add follower peer to followers.
store, ok := r.followers[storeID]
if !ok {
store = newRegionSubTree()
r.followers[storeID] = store
}
store.update(region)
}
}
// Add to learners.
for _, peer := range region.GetLearners() {
storeID := peer.GetStoreId()
store, ok := r.learners[storeID]
if !ok {
store = newRegionSubTree()
r.learners[storeID] = store
}
store.update(region)
}
for _, peer := range region.pendingPeers {
storeID := peer.GetStoreId()
store, ok := r.pendingPeers[storeID]
if !ok {
store = newRegionSubTree()
r.pendingPeers[storeID] = store
}
store.update(region)
}
return overlaps
}
// RemoveRegion removes RegionInfo from regionTree and regionMap
func (r *RegionsInfo) RemoveRegion(region *RegionInfo) {
// Remove from tree and regions.
r.tree.remove(region)
r.regions.Delete(region.GetID())
// Remove from leaders and followers.
for _, peer := range region.meta.GetPeers() {
storeID := peer.GetStoreId()
r.leaders[storeID].remove(region)
r.followers[storeID].remove(region)
r.learners[storeID].remove(region)
r.pendingPeers[storeID].remove(region)
}
}
// SearchRegion searches RegionInfo from regionTree
func (r *RegionsInfo) SearchRegion(regionKey []byte) *RegionInfo {
region := r.tree.search(regionKey)
if region == nil {
return nil
}
return r.GetRegion(region.GetID())
}
// SearchPrevRegion searches previous RegionInfo from regionTree
func (r *RegionsInfo) SearchPrevRegion(regionKey []byte) *RegionInfo {
region := r.tree.searchPrev(regionKey)
if region == nil {
return nil
}
return r.GetRegion(region.GetID())
}
// GetRegions gets all RegionInfo from regionMap
func (r *RegionsInfo) GetRegions() []*RegionInfo {
regions := make([]*RegionInfo, 0, r.regions.Len())
for _, region := range r.regions.m {
regions = append(regions, region)
}
return regions
}
// GetStoreRegions gets all RegionInfo with a given storeID
func (r *RegionsInfo) GetStoreRegions(storeID uint64) []*RegionInfo {
regions := make([]*RegionInfo, 0, r.GetStoreLeaderCount(storeID)+r.GetStoreFollowerCount(storeID))
if leaders, ok := r.leaders[storeID]; ok {
for _, region := range leaders.scanRanges() {
regions = append(regions, region)
}
}
if followers, ok := r.followers[storeID]; ok {
for _, region := range followers.scanRanges() {
regions = append(regions, region)
}
}
return regions
}
// GetStoreLeaderRegionSize gets total size of store's leader regions
func (r *RegionsInfo) GetStoreLeaderRegionSize(storeID uint64) int64 {
return r.leaders[storeID].TotalSize()
}
// GetStoreFollowerRegionSize gets total size of store's follower regions
func (r *RegionsInfo) GetStoreFollowerRegionSize(storeID uint64) int64 {
return r.followers[storeID].TotalSize()
}
// GetStoreLearnerRegionSize gets total size of store's learner regions
func (r *RegionsInfo) GetStoreLearnerRegionSize(storeID uint64) int64 {
return r.learners[storeID].TotalSize()
}
// GetStoreRegionSize gets total size of store's regions
func (r *RegionsInfo) GetStoreRegionSize(storeID uint64) int64 {
return r.GetStoreLeaderRegionSize(storeID) + r.GetStoreFollowerRegionSize(storeID) + r.GetStoreLearnerRegionSize(storeID)
}
// 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 {
regions = append(regions, proto.Clone(region.meta).(*metapb.Region))
}
return regions
}
// GetRegionCount gets the total count of RegionInfo of regionMap
func (r *RegionsInfo) GetRegionCount() int {
return r.regions.Len()
}
// GetStoreRegionCount gets the total count of a store's leader and follower RegionInfo by storeID
func (r *RegionsInfo) GetStoreRegionCount(storeID uint64) int {
return r.GetStoreLeaderCount(storeID) + r.GetStoreFollowerCount(storeID) + r.GetStoreLearnerCount(storeID)
}
// GetStorePendingPeerCount gets the total count of a store's region that includes pending peer
func (r *RegionsInfo) GetStorePendingPeerCount(storeID uint64) int {
return r.pendingPeers[storeID].length()
}
// GetStoreLeaderCount gets the total count of a store's leader RegionInfo
func (r *RegionsInfo) GetStoreLeaderCount(storeID uint64) int {
return r.leaders[storeID].length()
}
// GetStoreFollowerCount gets the total count of a store's follower RegionInfo
func (r *RegionsInfo) GetStoreFollowerCount(storeID uint64) int {
return r.followers[storeID].length()
}
// GetStoreLearnerCount gets the total count of a store's learner RegionInfo
func (r *RegionsInfo) GetStoreLearnerCount(storeID uint64) int {
return r.learners[storeID].length()
}
// RandRegion gets a region by random
func (r *RegionsInfo) RandRegion(opts ...RegionOption) *RegionInfo {
return randRegion(r.tree, opts...)
}
// RandPendingRegion randomly gets a store's region with a pending peer.
func (r *RegionsInfo) RandPendingRegion(storeID uint64, opts ...RegionOption) *RegionInfo {
return randRegion(r.pendingPeers[storeID], opts...)
}
// RandLeaderRegion randomly gets a store's leader region.
func (r *RegionsInfo) RandLeaderRegion(storeID uint64, opts ...RegionOption) *RegionInfo {
return randRegion(r.leaders[storeID], opts...)
}
// RandFollowerRegion randomly gets a store's follower region.
func (r *RegionsInfo) RandFollowerRegion(storeID uint64, opts ...RegionOption) *RegionInfo {
return randRegion(r.followers[storeID], opts...)
}
// GetPendingRegionsWithLock returns pending regions subtree by storeID
func (r *RegionsInfo) GetPendingRegionsWithLock(storeID uint64, callback func(RegionsContainer)) {
callback(r.pendingPeers[storeID])
}
// GetLeadersWithLock returns leaders subtree by storeID
func (r *RegionsInfo) GetLeadersWithLock(storeID uint64, callback func(RegionsContainer)) {
callback(r.leaders[storeID])
}
// GetFollowersWithLock returns followers subtree by storeID
func (r *RegionsInfo) GetFollowersWithLock(storeID uint64, callback func(RegionsContainer)) {
callback(r.followers[storeID])
}
// GetLeader returns leader RegionInfo by storeID and regionID(now only used in test)
func (r *RegionsInfo) GetLeader(storeID uint64, region *RegionInfo) *RegionInfo {
return r.leaders[storeID].find(region).region
}
// GetFollower returns follower RegionInfo by storeID and regionID(now only used in test)
func (r *RegionsInfo) GetFollower(storeID uint64, region *RegionInfo) *RegionInfo {
return r.followers[storeID].find(region).region
}
// ScanRange scans regions intersecting [start key, end key), returns at most
// `limit` regions. limit <= 0 means no limit.
func (r *RegionsInfo) ScanRange(startKey, endKey []byte, limit int) []*RegionInfo {
var res []*RegionInfo
r.tree.scanRange(startKey, func(region *RegionInfo) bool {
if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 {
return false
}
if limit > 0 && len(res) >= limit {
return false
}
res = append(res, r.GetRegion(region.GetID()))
return true
})
return res
}
// GetAverageRegionSize returns the average region approximate size.
func (r *RegionsInfo) GetAverageRegionSize() int64 {
if r.regions.Len() == 0 {
return 0
}
return r.regions.TotalSize() / int64(r.regions.Len())
}
const randomRegionMaxRetry = 10
// RegionsContainer is a container to store regions.
type RegionsContainer interface {
RandomRegion(startKey, endKey []byte) *RegionInfo
}
func randRegion(regions RegionsContainer, opts ...RegionOption) *RegionInfo {
for i := 0; i < randomRegionMaxRetry; i++ {
region := regions.RandomRegion(nil, nil)
if region == nil {
return nil
}
isSelect := true
for _, opt := range opts {
if !opt(region) {
isSelect = false
break
}
}
if isSelect {
return region
}
}
return nil
}
// DiffRegionPeersInfo returns the difference of peers info between two RegionInfo
func DiffRegionPeersInfo(origin *RegionInfo, other *RegionInfo) string {
var ret []string
for _, a := range origin.meta.Peers {
both := false
for _, b := range other.meta.Peers {
if reflect.DeepEqual(a, b) {
both = true
break
}
}
if !both {
ret = append(ret, fmt.Sprintf("Remove peer:{%v}", a))
}
}
for _, b := range other.meta.Peers {
both := false
for _, a := range origin.meta.Peers {
if reflect.DeepEqual(a, b) {
both = true
break
}
}
if !both {
ret = append(ret, fmt.Sprintf("Add peer:{%v}", b))
}
}
return strings.Join(ret, ",")
}
// DiffRegionKeyInfo returns the difference of key info between two RegionInfo
func DiffRegionKeyInfo(origin *RegionInfo, other *RegionInfo) string {
var ret []string
if !bytes.Equal(origin.meta.StartKey, other.meta.StartKey) {
ret = append(ret, fmt.Sprintf("StartKey Changed:{%s} -> {%s}", HexRegionKey(origin.meta.StartKey), HexRegionKey(other.meta.StartKey)))
} else {
ret = append(ret, fmt.Sprintf("StartKey:{%s}", HexRegionKey(origin.meta.StartKey)))
}
if !bytes.Equal(origin.meta.EndKey, other.meta.EndKey) {
ret = append(ret, fmt.Sprintf("EndKey Changed:{%s} -> {%s}", HexRegionKey(origin.meta.EndKey), HexRegionKey(other.meta.EndKey)))
} else {
ret = append(ret, fmt.Sprintf("EndKey:{%s}", HexRegionKey(origin.meta.EndKey)))
}
return strings.Join(ret, ", ")
}
// HexRegionKey converts region key to hex format. Used for formating region in
// logs.
func HexRegionKey(key []byte) []byte {
return []byte(strings.ToUpper(hex.EncodeToString(key)))
}
// RegionToHexMeta converts a region meta's keys to hex format. Used for formating
// region in logs.
func RegionToHexMeta(meta *metapb.Region) HexRegionMeta {
if meta == nil {
return HexRegionMeta{}
}
meta = proto.Clone(meta).(*metapb.Region)
meta.StartKey = HexRegionKey(meta.StartKey)
meta.EndKey = HexRegionKey(meta.EndKey)
return HexRegionMeta{meta}
}
// HexRegionMeta is a region meta in the hex format. Used for formating region in logs.
type HexRegionMeta struct {
*metapb.Region
}
func (h HexRegionMeta) String() string {
return strings.TrimSpace(proto.CompactTextString(h.Region))
}
// RegionsToHexMeta converts regions' meta keys to hex format. Used for formating
// region in logs.
func RegionsToHexMeta(regions []*metapb.Region) HexRegionsMeta {
hexRegionMetas := make([]*metapb.Region, len(regions))
for i, region := range regions {
meta := proto.Clone(region).(*metapb.Region)
meta.StartKey = HexRegionKey(meta.StartKey)
meta.EndKey = HexRegionKey(meta.EndKey)
hexRegionMetas[i] = meta
}
return HexRegionsMeta(hexRegionMetas)
}
// HexRegionsMeta is a slice of regions' meta in the hex format. Used for formating
// region in logs.
type HexRegionsMeta []*metapb.Region
func (h HexRegionsMeta) String() string {
var b strings.Builder
for _, r := range h {
b.WriteString(proto.CompactTextString(r))
}
return strings.TrimSpace(b.String())
}