cluster.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 mocktikv

import (
	"bytes"
	"context"
	"math"
	"sort"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/tablecodec"
	"go.uber.org/atomic"
)

// Cluster simulates a TiKV cluster. It focuses on management and the change of
// meta data. A Cluster mainly includes following 3 kinds of meta data:
// 1) Region: A Region is a fragment of TiKV's data whose range is [start, end).
//    The data of a Region is duplicated to multiple Peers and distributed in
//    multiple Stores.
// 2) Peer: A Peer is a replica of a Region's data. All peers of a Region form
//    a group, each group elects a Leader to provide services.
// 3) Store: A Store is a storage/service node. Try to think it as a TiKV server
//    process. Only the store with request's Region's leader Peer could respond
//    to client's request.
type Cluster struct {
	sync.RWMutex
	id      uint64
	stores  map[uint64]*Store
	regions map[uint64]*Region

	// delayEvents is used to control the execution sequence of rpc requests for test.
	delayEvents map[delayKey]time.Duration
	delayMu     sync.Mutex
}

type delayKey struct {
	startTS  uint64
	regionID uint64
}

// NewCluster creates an empty cluster. It needs to be bootstrapped before
// providing service.
func NewCluster() *Cluster {
	return &Cluster{
		stores:      make(map[uint64]*Store),
		regions:     make(map[uint64]*Region),
		delayEvents: make(map[delayKey]time.Duration),
	}
}

// AllocID creates an unique ID in cluster. The ID could be used as either
// StoreID, RegionID, or PeerID.
func (c *Cluster) AllocID() uint64 {
	c.Lock()
	defer c.Unlock()

	return c.allocID()
}

// AllocIDs creates multiple IDs.
func (c *Cluster) AllocIDs(n int) []uint64 {
	c.Lock()
	defer c.Unlock()

	var ids []uint64
	for len(ids) < n {
		ids = append(ids, c.allocID())
	}
	return ids
}

func (c *Cluster) allocID() uint64 {
	c.id++
	return c.id
}

// GetAllRegions gets all the regions in the cluster.
func (c *Cluster) GetAllRegions() []*Region {
	regions := make([]*Region, 0, len(c.regions))
	for _, region := range c.regions {
		regions = append(regions, region)
	}
	return regions
}

// GetStore returns a Store's meta.
func (c *Cluster) GetStore(storeID uint64) *metapb.Store {
	c.RLock()
	defer c.RUnlock()

	if store := c.stores[storeID]; store != nil {
		return proto.Clone(store.meta).(*metapb.Store)
	}
	return nil
}

// GetAllStores returns all Stores' meta.
func (c *Cluster) GetAllStores() []*metapb.Store {
	c.RLock()
	defer c.RUnlock()

	stores := make([]*metapb.Store, 0, len(c.stores))
	for _, store := range c.stores {
		stores = append(stores, proto.Clone(store.meta).(*metapb.Store))
	}
	return stores
}

// StopStore stops a store with storeID.
func (c *Cluster) StopStore(storeID uint64) {
	c.Lock()
	defer c.Unlock()

	if store := c.stores[storeID]; store != nil {
		store.meta.State = metapb.StoreState_Offline
	}
}

// StartStore starts a store with storeID.
func (c *Cluster) StartStore(storeID uint64) {
	c.Lock()
	defer c.Unlock()

	if store := c.stores[storeID]; store != nil {
		store.meta.State = metapb.StoreState_Up
	}
}

// CancelStore makes the store with cancel state true.
func (c *Cluster) CancelStore(storeID uint64) {
	c.Lock()
	defer c.Unlock()

	//A store returns context.Cancelled Error when cancel is true.
	if store := c.stores[storeID]; store != nil {
		store.cancel = true
	}
}

// UnCancelStore makes the store with cancel state false.
func (c *Cluster) UnCancelStore(storeID uint64) {
	c.Lock()
	defer c.Unlock()

	if store := c.stores[storeID]; store != nil {
		store.cancel = false
	}
}

// GetStoreByAddr returns a Store's meta by an addr.
func (c *Cluster) GetStoreByAddr(addr string) *metapb.Store {
	c.RLock()
	defer c.RUnlock()

	for _, s := range c.stores {
		if s.meta.GetAddress() == addr {
			return proto.Clone(s.meta).(*metapb.Store)
		}
	}
	return nil
}

// GetAndCheckStoreByAddr checks and returns a Store's meta by an addr
func (c *Cluster) GetAndCheckStoreByAddr(addr string) (*metapb.Store, error) {
	c.RLock()
	defer c.RUnlock()

	for _, s := range c.stores {
		if s.cancel {
			return nil, context.Canceled
		}
		if s.meta.GetAddress() == addr {
			return proto.Clone(s.meta).(*metapb.Store), nil
		}
	}
	return nil, nil
}

// AddStore add a new Store to the cluster.
func (c *Cluster) AddStore(storeID uint64, addr string) {
	c.Lock()
	defer c.Unlock()

	c.stores[storeID] = newStore(storeID, addr)
}

// RemoveStore removes a Store from the cluster.
func (c *Cluster) RemoveStore(storeID uint64) {
	c.Lock()
	defer c.Unlock()

	delete(c.stores, storeID)
}

// UpdateStoreAddr updates store address for cluster.
func (c *Cluster) UpdateStoreAddr(storeID uint64, addr string) {
	c.Lock()
	defer c.Unlock()
	c.stores[storeID] = newStore(storeID, addr)
}

// GetRegion returns a Region's meta and leader ID.
func (c *Cluster) GetRegion(regionID uint64) (*metapb.Region, uint64) {
	c.RLock()
	defer c.RUnlock()

	r := c.regions[regionID]
	if r == nil {
		return nil, 0
	}
	return proto.Clone(r.Meta).(*metapb.Region), r.leader
}

// GetRegionByKey returns the Region and its leader whose range contains the key.
func (c *Cluster) GetRegionByKey(key []byte) (*metapb.Region, *metapb.Peer) {
	c.RLock()
	defer c.RUnlock()

	for _, r := range c.regions {
		if regionContains(r.Meta.StartKey, r.Meta.EndKey, key) {
			return proto.Clone(r.Meta).(*metapb.Region), proto.Clone(r.leaderPeer()).(*metapb.Peer)
		}
	}
	return nil, nil
}

// GetPrevRegionByKey returns the previous Region and its leader whose range contains the key.
func (c *Cluster) GetPrevRegionByKey(key []byte) (*metapb.Region, *metapb.Peer) {
	c.RLock()
	defer c.RUnlock()

	currentRegion, _ := c.GetRegionByKey(key)
	if len(currentRegion.StartKey) == 0 {
		return nil, nil
	}
	for _, r := range c.regions {
		if bytes.Equal(r.Meta.EndKey, currentRegion.StartKey) {
			return proto.Clone(r.Meta).(*metapb.Region), proto.Clone(r.leaderPeer()).(*metapb.Peer)
		}
	}
	return nil, nil
}

// GetRegionByID returns the Region and its leader whose ID is regionID.
func (c *Cluster) GetRegionByID(regionID uint64) (*metapb.Region, *metapb.Peer) {
	c.RLock()
	defer c.RUnlock()

	for _, r := range c.regions {
		if r.Meta.GetId() == regionID {
			return proto.Clone(r.Meta).(*metapb.Region), proto.Clone(r.leaderPeer()).(*metapb.Peer)
		}
	}
	return nil, nil
}

// ScanRegions returns at most `limit` regions from given `key` and their leaders.
func (c *Cluster) ScanRegions(startKey, endKey []byte, limit int) ([]*metapb.Region, []*metapb.Peer) {
	c.RLock()
	defer c.RUnlock()

	regions := make([]*Region, 0, len(c.regions))
	for _, region := range c.regions {
		regions = append(regions, region)
	}

	sort.Slice(regions, func(i, j int) bool {
		return bytes.Compare(regions[i].Meta.GetStartKey(), regions[j].Meta.GetStartKey()) < 0
	})

	startPos := sort.Search(len(regions), func(i int) bool {
		if len(regions[i].Meta.GetEndKey()) == 0 {
			return true
		}
		return bytes.Compare(regions[i].Meta.GetEndKey(), startKey) > 0
	})
	regions = regions[startPos:]
	if len(endKey) > 0 {
		endPos := sort.Search(len(regions), func(i int) bool {
			return bytes.Compare(regions[i].Meta.GetStartKey(), endKey) >= 0
		})
		if endPos > 0 {
			regions = regions[:endPos]
		}
	}
	if limit > 0 && len(regions) > limit {
		regions = regions[:limit]
	}

	metas := make([]*metapb.Region, 0, len(regions))
	leaders := make([]*metapb.Peer, 0, len(regions))
	for _, region := range regions {
		leader := region.leaderPeer()
		if leader == nil {
			leader = &metapb.Peer{}
		} else {
			leader = proto.Clone(leader).(*metapb.Peer)
		}

		metas = append(metas, proto.Clone(region.Meta).(*metapb.Region))
		leaders = append(leaders, leader)
	}

	return metas, leaders
}

// Bootstrap creates the first Region. The Stores should be in the Cluster before
// bootstrap.
func (c *Cluster) Bootstrap(regionID uint64, storeIDs, peerIDs []uint64, leaderPeerID uint64) {
	c.Lock()
	defer c.Unlock()

	if len(storeIDs) != len(peerIDs) {
		panic("len(storeIDs) != len(peerIDs)")
	}
	c.regions[regionID] = newRegion(regionID, storeIDs, peerIDs, leaderPeerID)
}

// AddPeer adds a new Peer for the Region on the Store.
func (c *Cluster) AddPeer(regionID, storeID, peerID uint64) {
	c.Lock()
	defer c.Unlock()

	c.regions[regionID].addPeer(peerID, storeID)
}

// RemovePeer removes the Peer from the Region. Note that if the Peer is leader,
// the Region will have no leader before calling ChangeLeader().
func (c *Cluster) RemovePeer(regionID, storeID uint64) {
	c.Lock()
	defer c.Unlock()

	c.regions[regionID].removePeer(storeID)
}

// ChangeLeader sets the Region's leader Peer. Caller should guarantee the Peer
// exists.
func (c *Cluster) ChangeLeader(regionID, leaderPeerID uint64) {
	c.Lock()
	defer c.Unlock()

	c.regions[regionID].changeLeader(leaderPeerID)
}

// GiveUpLeader sets the Region's leader to 0. The Region will have no leader
// before calling ChangeLeader().
func (c *Cluster) GiveUpLeader(regionID uint64) {
	c.ChangeLeader(regionID, 0)
}

// Split splits a Region at the key (encoded) and creates new Region.
func (c *Cluster) Split(regionID, newRegionID uint64, key []byte, peerIDs []uint64, leaderPeerID uint64) {
	c.SplitRaw(regionID, newRegionID, NewMvccKey(key), peerIDs, leaderPeerID)
}

// SplitRaw splits a Region at the key (not encoded) and creates new Region.
func (c *Cluster) SplitRaw(regionID, newRegionID uint64, rawKey []byte, peerIDs []uint64, leaderPeerID uint64) *metapb.Region {
	c.Lock()
	defer c.Unlock()

	newRegion := c.regions[regionID].split(newRegionID, rawKey, peerIDs, leaderPeerID)
	c.regions[newRegionID] = newRegion
	// The mocktikv should return a deep copy of meta info to avoid data race
	meta := proto.Clone(newRegion.Meta)
	return meta.(*metapb.Region)
}

// Merge merges 2 regions, their key ranges should be adjacent.
func (c *Cluster) Merge(regionID1, regionID2 uint64) {
	c.Lock()
	defer c.Unlock()

	c.regions[regionID1].merge(c.regions[regionID2].Meta.GetEndKey())
	delete(c.regions, regionID2)
}

// SplitTable evenly splits the data in table into count regions.
// Only works for single store.
func (c *Cluster) SplitTable(mvccStore MVCCStore, tableID int64, count int) {
	tableStart := tablecodec.GenTableRecordPrefix(tableID)
	tableEnd := tableStart.PrefixNext()
	c.splitRange(mvccStore, NewMvccKey(tableStart), NewMvccKey(tableEnd), count)
}

// SplitIndex evenly splits the data in index into count regions.
// Only works for single store.
func (c *Cluster) SplitIndex(mvccStore MVCCStore, tableID, indexID int64, count int) {
	indexStart := tablecodec.EncodeTableIndexPrefix(tableID, indexID)
	indexEnd := indexStart.PrefixNext()
	c.splitRange(mvccStore, NewMvccKey(indexStart), NewMvccKey(indexEnd), count)
}

// SplitKeys evenly splits the start, end key into "count" regions.
// Only works for single store.
func (c *Cluster) SplitKeys(mvccStore MVCCStore, start, end kv.Key, count int) {
	c.splitRange(mvccStore, NewMvccKey(start), NewMvccKey(end), count)
}

// ScheduleDelay schedules a delay event for a transaction on a region.
func (c *Cluster) ScheduleDelay(startTS, regionID uint64, dur time.Duration) {
	c.delayMu.Lock()
	c.delayEvents[delayKey{startTS: startTS, regionID: regionID}] = dur
	c.delayMu.Unlock()
}

func (c *Cluster) handleDelay(startTS, regionID uint64) {
	key := delayKey{startTS: startTS, regionID: regionID}
	c.delayMu.Lock()
	dur, ok := c.delayEvents[key]
	if ok {
		delete(c.delayEvents, key)
	}
	c.delayMu.Unlock()
	if ok {
		time.Sleep(dur)
	}
}

func (c *Cluster) splitRange(mvccStore MVCCStore, start, end MvccKey, count int) {
	c.Lock()
	defer c.Unlock()
	c.evacuateOldRegionRanges(start, end)
	regionPairs := c.getEntriesGroupByRegions(mvccStore, start, end, count)
	c.createNewRegions(regionPairs, start, end)
}

// getEntriesGroupByRegions groups the key value pairs into splitted regions.
func (c *Cluster) getEntriesGroupByRegions(mvccStore MVCCStore, start, end MvccKey, count int) [][]Pair {
	startTS := uint64(math.MaxUint64)
	limit := math.MaxInt32
	pairs := mvccStore.Scan(start.Raw(), end.Raw(), limit, startTS, kvrpcpb.IsolationLevel_SI, nil)
	regionEntriesSlice := make([][]Pair, 0, count)
	quotient := len(pairs) / count
	remainder := len(pairs) % count
	i := 0
	for i < len(pairs) {
		regionEntryCount := quotient
		if remainder > 0 {
			remainder--
			regionEntryCount++
		}
		regionEntries := pairs[i : i+regionEntryCount]
		regionEntriesSlice = append(regionEntriesSlice, regionEntries)
		i += regionEntryCount
	}
	return regionEntriesSlice
}

func (c *Cluster) createNewRegions(regionPairs [][]Pair, start, end MvccKey) {
	for i := range regionPairs {
		peerID := c.allocID()
		newRegion := newRegion(c.allocID(), []uint64{c.firstStoreID()}, []uint64{peerID}, peerID)
		var regionStartKey, regionEndKey MvccKey
		if i == 0 {
			regionStartKey = start
		} else {
			regionStartKey = NewMvccKey(regionPairs[i][0].Key)
		}
		if i == len(regionPairs)-1 {
			regionEndKey = end
		} else {
			// Use the next region's first key as region end key.
			regionEndKey = NewMvccKey(regionPairs[i+1][0].Key)
		}
		newRegion.updateKeyRange(regionStartKey, regionEndKey)
		c.regions[newRegion.Meta.Id] = newRegion
	}
}

// evacuateOldRegionRanges evacuate the range [start, end].
// Old regions has intersection with [start, end) will be updated or deleted.
func (c *Cluster) evacuateOldRegionRanges(start, end MvccKey) {
	oldRegions := c.getRegionsCoverRange(start, end)
	for _, oldRegion := range oldRegions {
		startCmp := bytes.Compare(oldRegion.Meta.StartKey, start)
		endCmp := bytes.Compare(oldRegion.Meta.EndKey, end)
		if len(oldRegion.Meta.EndKey) == 0 {
			endCmp = 1
		}
		if startCmp >= 0 && endCmp <= 0 {
			// The region is within table data, it will be replaced by new regions.
			delete(c.regions, oldRegion.Meta.Id)
		} else if startCmp < 0 && endCmp > 0 {
			// A single Region covers table data, split into two regions that do not overlap table data.
			oldEnd := oldRegion.Meta.EndKey
			oldRegion.updateKeyRange(oldRegion.Meta.StartKey, start)
			peerID := c.allocID()
			newRegion := newRegion(c.allocID(), []uint64{c.firstStoreID()}, []uint64{peerID}, peerID)
			newRegion.updateKeyRange(end, oldEnd)
			c.regions[newRegion.Meta.Id] = newRegion
		} else if startCmp < 0 {
			oldRegion.updateKeyRange(oldRegion.Meta.StartKey, start)
		} else {
			oldRegion.updateKeyRange(end, oldRegion.Meta.EndKey)
		}
	}
}

func (c *Cluster) firstStoreID() uint64 {
	for id := range c.stores {
		return id
	}
	return 0
}

// getRegionsCoverRange gets regions in the cluster that has intersection with [start, end).
func (c *Cluster) getRegionsCoverRange(start, end MvccKey) []*Region {
	regions := make([]*Region, 0, len(c.regions))
	for _, region := range c.regions {
		onRight := bytes.Compare(end, region.Meta.StartKey) <= 0
		onLeft := bytes.Compare(region.Meta.EndKey, start) <= 0
		if len(region.Meta.EndKey) == 0 {
			onLeft = false
		}
		if onLeft || onRight {
			continue
		}
		regions = append(regions, region)
	}
	return regions
}

// Region is the Region meta data.
type Region struct {
	Meta   *metapb.Region
	leader uint64
}

func newPeerMeta(peerID, storeID uint64) *metapb.Peer {
	return &metapb.Peer{
		Id:      peerID,
		StoreId: storeID,
	}
}

func newRegion(regionID uint64, storeIDs, peerIDs []uint64, leaderPeerID uint64) *Region {
	if len(storeIDs) != len(peerIDs) {
		panic("len(storeIDs) != len(peerIds)")
	}
	peers := make([]*metapb.Peer, 0, len(storeIDs))
	for i := range storeIDs {
		peers = append(peers, newPeerMeta(peerIDs[i], storeIDs[i]))
	}
	meta := &metapb.Region{
		Id:    regionID,
		Peers: peers,
	}
	return &Region{
		Meta:   meta,
		leader: leaderPeerID,
	}
}

func (r *Region) addPeer(peerID, storeID uint64) {
	r.Meta.Peers = append(r.Meta.Peers, newPeerMeta(peerID, storeID))
	r.incConfVer()
}

func (r *Region) removePeer(peerID uint64) {
	for i, peer := range r.Meta.Peers {
		if peer.GetId() == peerID {
			r.Meta.Peers = append(r.Meta.Peers[:i], r.Meta.Peers[i+1:]...)
			break
		}
	}
	if r.leader == peerID {
		r.leader = 0
	}
	r.incConfVer()
}

func (r *Region) changeLeader(leaderID uint64) {
	r.leader = leaderID
}

func (r *Region) leaderPeer() *metapb.Peer {
	for _, p := range r.Meta.Peers {
		if p.GetId() == r.leader {
			return p
		}
	}
	return nil
}

func (r *Region) split(newRegionID uint64, key MvccKey, peerIDs []uint64, leaderPeerID uint64) *Region {
	if len(r.Meta.Peers) != len(peerIDs) {
		panic("len(r.meta.Peers) != len(peerIDs)")
	}
	storeIDs := make([]uint64, 0, len(r.Meta.Peers))
	for _, peer := range r.Meta.Peers {
		storeIDs = append(storeIDs, peer.GetStoreId())
	}
	region := newRegion(newRegionID, storeIDs, peerIDs, leaderPeerID)
	region.updateKeyRange(key, r.Meta.EndKey)
	r.updateKeyRange(r.Meta.StartKey, key)
	return region
}

func (r *Region) merge(endKey MvccKey) {
	r.Meta.EndKey = endKey
	r.incVersion()
}

func (r *Region) updateKeyRange(start, end MvccKey) {
	r.Meta.StartKey = start
	r.Meta.EndKey = end
	r.incVersion()
}

func (r *Region) incConfVer() {
	r.Meta.RegionEpoch = &metapb.RegionEpoch{
		ConfVer: r.Meta.GetRegionEpoch().GetConfVer() + 1,
		Version: r.Meta.GetRegionEpoch().GetVersion(),
	}
}

func (r *Region) incVersion() {
	r.Meta.RegionEpoch = &metapb.RegionEpoch{
		ConfVer: r.Meta.GetRegionEpoch().GetConfVer(),
		Version: r.Meta.GetRegionEpoch().GetVersion() + 1,
	}
}

// Store is the Store's meta data.
type Store struct {
	meta       *metapb.Store
	cancel     bool // return context.Cancelled error when cancel is true.
	tokenCount atomic.Int64
}

func newStore(storeID uint64, addr string) *Store {
	return &Store{
		meta: &metapb.Store{
			Id:      storeID,
			Address: addr,
		},
	}
}