func (mc *Cluster) UpdateStorageWrittenStats(storeID, bytesWritten, keysWritten uint64) {
	mc.updateStorageStatistics(storeID, func(newStats *pdpb.StoreStats) {
		newStats.BytesWritten = bytesWritten
		newStats.KeysWritten = keysWritten
	})
}
UpdateStorageWrittenStats updates store written bytes.
package mockcluster
import (
	"context"
	"fmt"
	"strconv"
	"time"
	"github.com/gogo/protobuf/proto"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/log"
	"github.com/tikv/pd/pkg/errs"
	"github.com/tikv/pd/pkg/mock/mockid"
	"github.com/tikv/pd/server/config"
	"github.com/tikv/pd/server/core"
	"github.com/tikv/pd/server/core/storelimit"
	"github.com/tikv/pd/server/kv"
	"github.com/tikv/pd/server/schedule/labeler"
	"github.com/tikv/pd/server/schedule/placement"
	"github.com/tikv/pd/server/statistics"
	"github.com/tikv/pd/server/versioninfo"
)
const (
	defaultStoreCapacity = 100 * (1 << 30)  
	defaultRegionSize    = 96 * (1 << 20)   
	mb                   = (1 << 20)        
)
type Cluster struct {
	*core.BasicCluster
	*mockid.IDAllocator
	*placement.RuleManager
	*labeler.RegionLabeler
	*statistics.HotStat
	*config.PersistOptions
	ID               uint64
	suspectRegions   map[uint64]struct{}
	disabledFeatures map[versioninfo.Feature]struct{}
}
func NewCluster(ctx context.Context, opts *config.PersistOptions) *Cluster {
	mockQuit := make(chan struct{})
	clus := &Cluster{
		BasicCluster:     core.NewBasicCluster(),
		IDAllocator:      mockid.NewIDAllocator(),
		HotStat:          statistics.NewHotStat(ctx, mockQuit),
		PersistOptions:   opts,
		suspectRegions:   map[uint64]struct{}{},
		disabledFeatures: make(map[versioninfo.Feature]struct{}),
	}
	if clus.PersistOptions.GetReplicationConfig().EnablePlacementRules {
		clus.initRuleManager()
	}
	clus.RegionLabeler, _ = labeler.NewRegionLabeler(core.NewStorage(kv.NewMemoryKV()))
	return clus
}
func (mc *Cluster) GetOpts() *config.PersistOptions {
	return mc.PersistOptions
}
func (mc *Cluster) AllocID() (uint64, error) {
	return mc.Alloc()
}
func (mc *Cluster) ScanRegions(startKey, endKey []byte, limit int) []*core.RegionInfo {
	return mc.Regions.ScanRange(startKey, endKey, limit)
}
func (mc *Cluster) LoadRegion(regionID uint64, peerStoreIDs ...uint64) {
	r := mc.newMockRegionInfo(regionID, 0, peerStoreIDs...).Clone(core.WithLeader(nil))
	mc.PutRegion(r)
}
func (mc *Cluster) GetStoresLoads() map[uint64][]float64 {
	mc.HotStat.FilterUnhealthyStore(mc)
	return mc.HotStat.GetStoresLoads()
}
func (mc *Cluster) GetStoreRegionCount(storeID uint64) int {
	return mc.Regions.GetStoreRegionCount(storeID)
}
func (mc *Cluster) GetStore(storeID uint64) *core.StoreInfo {
	return mc.Stores.GetStore(storeID)
}
func (mc *Cluster) IsRegionHot(region *core.RegionInfo) bool {
	return mc.HotCache.IsRegionHot(region, mc.GetHotRegionCacheHitsThreshold())
}
func (mc *Cluster) RegionReadStats() map[uint64][]*statistics.HotPeerStat {
	return mc.HotCache.RegionStats(statistics.ReadFlow, mc.GetHotRegionCacheHitsThreshold())
}
func (mc *Cluster) RegionWriteStats() map[uint64][]*statistics.HotPeerStat {
	return mc.HotCache.RegionStats(statistics.WriteFlow, mc.GetHotRegionCacheHitsThreshold())
}
func (mc *Cluster) HotRegionsFromStore(store uint64, kind statistics.FlowKind) []*core.RegionInfo {
	stats := mc.HotCache.HotRegionsFromStore(store, kind, mc.GetHotRegionCacheHitsThreshold())
	regions := make([]*core.RegionInfo, 0, len(stats))
	for _, stat := range stats {
		region := mc.GetRegion(stat.RegionID)
		if region != nil {
			regions = append(regions, region)
		}
	}
	return regions
}
func (mc *Cluster) AllocPeer(storeID uint64) (*metapb.Peer, error) {
	peerID, err := mc.AllocID()
	if err != nil {
		log.Error("failed to alloc peer", errs.ZapError(err))
		return nil, err
	}
	peer := &metapb.Peer{
		Id:      peerID,
		StoreId: storeID,
	}
	return peer, nil
}
func (mc *Cluster) initRuleManager() {
	if mc.RuleManager == nil {
		mc.RuleManager = placement.NewRuleManager(core.NewStorage(kv.NewMemoryKV()), mc, mc.GetOpts())
		mc.RuleManager.Initialize(int(mc.GetReplicationConfig().MaxReplicas), mc.GetReplicationConfig().LocationLabels)
	}
}
func (mc *Cluster) GetRuleManager() *placement.RuleManager {
	return mc.RuleManager
}
func (mc *Cluster) GetRegionLabeler() *labeler.RegionLabeler {
	return mc.RegionLabeler
}
func (mc *Cluster) SetStoreUp(storeID uint64) {
	store := mc.GetStore(storeID)
	newStore := store.Clone(
		core.UpStore(),
		core.SetLastHeartbeatTS(time.Now()),
	)
	mc.PutStore(newStore)
}
func (mc *Cluster) SetStoreDisconnect(storeID uint64) {
	store := mc.GetStore(storeID)
	newStore := store.Clone(
		core.UpStore(),
		core.SetLastHeartbeatTS(time.Now().Add(-time.Second*30)),
	)
	mc.PutStore(newStore)
}
func (mc *Cluster) SetStoreDown(storeID uint64) {
	store := mc.GetStore(storeID)
	newStore := store.Clone(
		core.UpStore(),
		core.SetLastHeartbeatTS(time.Time{}),
	)
	mc.PutStore(newStore)
}
func (mc *Cluster) SetStoreOffline(storeID uint64) {
	store := mc.GetStore(storeID)
	newStore := store.Clone(core.OfflineStore(false))
	mc.PutStore(newStore)
}
func (mc *Cluster) SetStoreBusy(storeID uint64, busy bool) {
	store := mc.GetStore(storeID)
	newStats := proto.Clone(store.GetStoreStats()).(*pdpb.StoreStats)
	newStats.IsBusy = busy
	newStore := store.Clone(
		core.SetStoreStats(newStats),
		core.SetLastHeartbeatTS(time.Now()),
	)
	mc.PutStore(newStore)
}
func (mc *Cluster) AddLeaderStore(storeID uint64, leaderCount int, leaderSizes ...int64) {
	stats := &pdpb.StoreStats{}
	stats.Capacity = defaultStoreCapacity
	stats.UsedSize = uint64(leaderCount) * defaultRegionSize
	stats.Available = stats.Capacity - uint64(leaderCount)*defaultRegionSize
	var leaderSize int64
	if len(leaderSizes) != 0 {
		leaderSize = leaderSizes[0]
	} else {
		leaderSize = int64(leaderCount) * defaultRegionSize / mb
	}
	store := core.NewStoreInfo(
		&metapb.Store{Id: storeID},
		core.SetStoreStats(stats),
		core.SetLeaderCount(leaderCount),
		core.SetLeaderSize(leaderSize),
		core.SetLastHeartbeatTS(time.Now()),
	)
	mc.SetStoreLimit(storeID, storelimit.AddPeer, 60)
	mc.SetStoreLimit(storeID, storelimit.RemovePeer, 60)
	mc.PutStore(store)
}
func (mc *Cluster) AddRegionStore(storeID uint64, regionCount int) {
	stats := &pdpb.StoreStats{}
	stats.Capacity = defaultStoreCapacity
	stats.UsedSize = uint64(regionCount) * defaultRegionSize
	stats.Available = stats.Capacity - uint64(regionCount)*defaultRegionSize
	store := core.NewStoreInfo(
		&metapb.Store{Id: storeID, Labels: []*metapb.StoreLabel{
			{
				Key:   "ID",
				Value: fmt.Sprintf("%v", storeID),
			},
		}},
		core.SetStoreStats(stats),
		core.SetRegionCount(regionCount),
		core.SetRegionSize(int64(regionCount)*defaultRegionSize/mb),
		core.SetLastHeartbeatTS(time.Now()),
	)
	mc.SetStoreLimit(storeID, storelimit.AddPeer, 60)
	mc.SetStoreLimit(storeID, storelimit.RemovePeer, 60)
	mc.PutStore(store)
}
func (mc *Cluster) AddRegionStoreWithLeader(storeID uint64, regionCount int, leaderCounts ...int) {
	leaderCount := regionCount
	if len(leaderCounts) != 0 {
		leaderCount = leaderCounts[0]
	}
	mc.AddRegionStore(storeID, regionCount)
	for i := 0; i < leaderCount; i++ {
		id, _ := mc.AllocID()
		mc.AddLeaderRegion(id, storeID)
	}
}
func (mc *Cluster) AddLabelsStore(storeID uint64, regionCount int, labels map[string]string) {
	newLabels := make([]*metapb.StoreLabel, 0, len(labels))
	for k, v := range labels {
		newLabels = append(newLabels, &metapb.StoreLabel{Key: k, Value: v})
	}
	stats := &pdpb.StoreStats{}
	stats.Capacity = defaultStoreCapacity
	stats.Available = stats.Capacity - uint64(regionCount)*defaultRegionSize
	stats.UsedSize = uint64(regionCount) * defaultRegionSize
	store := core.NewStoreInfo(
		&metapb.Store{
			Id:     storeID,
			Labels: newLabels,
		},
		core.SetStoreStats(stats),
		core.SetRegionCount(regionCount),
		core.SetRegionSize(int64(regionCount)*defaultRegionSize/mb),
		core.SetLastHeartbeatTS(time.Now()),
	)
	mc.SetStoreLimit(storeID, storelimit.AddPeer, 60)
	mc.SetStoreLimit(storeID, storelimit.RemovePeer, 60)
	mc.PutStore(store)
}
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
}
func (mc *Cluster) AddRegionWithLearner(regionID uint64, leaderStoreID uint64, followerStoreIDs, learnerStoreIDs []uint64) *core.RegionInfo {
	origin := mc.MockRegionInfo(regionID, leaderStoreID, followerStoreIDs, learnerStoreIDs, nil)
	region := origin.Clone(core.SetApproximateSize(defaultRegionSize/mb), core.SetApproximateKeys(10))
	mc.PutRegion(region)
	return region
}
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)),
	)
	mc.PutRegion(r)
}
func (mc *Cluster) AddRegionWithReadInfo(
	regionID uint64, leaderStoreID uint64,
	readBytes, readKeys, readQuery uint64,
	reportInterval uint64,
	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))
	r = r.Clone(core.SetReadQuery(readQuery))
	filledNum := mc.HotCache.GetFilledPeriod(statistics.ReadFlow)
	if len(filledNums) > 0 {
		filledNum = filledNums[0]
	}
	var items []*statistics.HotPeerStat
	for i := 0; i < filledNum; i++ {
		items = mc.CheckRegionRead(r)
		for _, item := range items {
			mc.HotCache.Update(item)
		}
	}
	mc.PutRegion(r)
	return items
}
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 {
		filledNum = filledNums[0]
	}
	var items []*statistics.HotPeerStat
	for i := 0; i < filledNum; i++ {
		items = mc.CheckRegionRead(r)
		for _, item := range items {
			if item.StoreID == targetStoreID {
				mc.HotCache.Update(item)
			}
		}
	}
	mc.PutRegion(r)
	return items
}
func (mc *Cluster) AddRegionLeaderWithReadInfo(
	regionID uint64, leaderStoreID uint64,
	readBytes, readKeys, readQuery uint64,
	reportInterval uint64,
	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.SetReadQuery(readQuery))
	r = r.Clone(core.SetReportInterval(reportInterval))
	filledNum := mc.HotCache.GetFilledPeriod(statistics.ReadFlow)
	if len(filledNums) > 0 {
		filledNum = filledNums[0]
	}
	var items []*statistics.HotPeerStat
	for i := 0; i < filledNum; i++ {
		items = mc.CheckRegionLeaderRead(r)
		for _, item := range items {
			mc.HotCache.Update(item)
		}
	}
	mc.PutRegion(r)
	return items
}
func (mc *Cluster) AddLeaderRegionWithWriteInfo(
	regionID uint64, leaderStoreID uint64,
	writtenBytes, writtenKeys, writtenQuery uint64,
	reportInterval uint64,
	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))
	r = r.Clone(core.SetWrittenQuery(writtenQuery))
	filledNum := mc.HotCache.GetFilledPeriod(statistics.WriteFlow)
	if len(filledNums) > 0 {
		filledNum = filledNums[0]
	}
	var items []*statistics.HotPeerStat
	for i := 0; i < filledNum; i++ {
		items = mc.CheckRegionWrite(r)
		for _, item := range items {
			mc.HotCache.Update(item)
		}
	}
	mc.PutRegion(r)
	return items
}
func (mc *Cluster) UpdateStoreLeaderWeight(storeID uint64, weight float64) {
	store := mc.GetStore(storeID)
	newStore := store.Clone(core.SetLeaderWeight(weight))
	mc.PutStore(newStore)
}
func (mc *Cluster) SetStoreEvictLeader(storeID uint64, enableEvictLeader bool) {
	store := mc.GetStore(storeID)
	if enableEvictLeader {
		mc.PutStore(store.Clone(core.PauseLeaderTransfer()))
	} else {
		mc.PutStore(store.Clone(core.ResumeLeaderTransfer()))
	}
}
func (mc *Cluster) UpdateStoreRegionWeight(storeID uint64, weight float64) {
	store := mc.GetStore(storeID)
	newStore := store.Clone(core.SetRegionWeight(weight))
	mc.PutStore(newStore)
}
func (mc *Cluster) UpdateStoreLeaderSize(storeID uint64, size int64) {
	store := mc.GetStore(storeID)
	newStats := proto.Clone(store.GetStoreStats()).(*pdpb.StoreStats)
	newStats.Available = newStats.Capacity - uint64(store.GetLeaderSize())
	newStore := store.Clone(
		core.SetStoreStats(newStats),
		core.SetLeaderSize(size),
	)
	mc.PutStore(newStore)
}
func (mc *Cluster) UpdateStoreRegionSize(storeID uint64, size int64) {
	store := mc.GetStore(storeID)
	newStats := proto.Clone(store.GetStoreStats()).(*pdpb.StoreStats)
	newStats.Available = newStats.Capacity - uint64(store.GetRegionSize())
	newStore := store.Clone(
		core.SetStoreStats(newStats),
		core.SetRegionSize(size),
	)
	mc.PutStore(newStore)
}
func (mc *Cluster) UpdateLeaderCount(storeID uint64, leaderCount int) {
	store := mc.GetStore(storeID)
	newStore := store.Clone(
		core.SetLeaderCount(leaderCount),
		core.SetLeaderSize(int64(leaderCount)*defaultRegionSize/mb),
	)
	mc.PutStore(newStore)
}
func (mc *Cluster) UpdateRegionCount(storeID uint64, regionCount int) {
	store := mc.GetStore(storeID)
	newStore := store.Clone(
		core.SetRegionCount(regionCount),
		core.SetRegionSize(int64(regionCount)*defaultRegionSize/mb),
	)
	mc.PutStore(newStore)
}
func (mc *Cluster) UpdateSnapshotCount(storeID uint64, snapshotCount int) {
	mc.updateStorageStatistics(storeID, func(newStats *pdpb.StoreStats) {
		newStats.ReceivingSnapCount = uint32(snapshotCount)
	})
}
func (mc *Cluster) UpdatePendingPeerCount(storeID uint64, pendingPeerCount int) {
	store := mc.GetStore(storeID)
	newStore := store.Clone(core.SetPendingPeerCount(pendingPeerCount))
	mc.PutStore(newStore)
}
func (mc *Cluster) UpdateStorageRatio(storeID uint64, usedRatio, availableRatio float64) {
	mc.updateStorageStatistics(storeID, func(newStats *pdpb.StoreStats) {
		newStats.Capacity = defaultStoreCapacity
		newStats.UsedSize = uint64(float64(newStats.Capacity) * usedRatio)
		newStats.Available = uint64(float64(newStats.Capacity) * availableRatio)
	})
}