talent-plan-tinykv/kv/test_raftstore/utils.go

package test_raftstore

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"time"

	"github.com/Connor1996/badger"
	"github.com/pingcap-incubator/tinykv/kv/config"
	"github.com/pingcap-incubator/tinykv/kv/util/engine_util"
	"github.com/pingcap-incubator/tinykv/log"
	"github.com/pingcap-incubator/tinykv/proto/pkg/metapb"
	"github.com/pingcap-incubator/tinykv/proto/pkg/raft_cmdpb"
)

func SleepMS(ms int64) {
	time.Sleep(time.Duration(ms) * time.Millisecond)
}

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

func NewBaseRequest(regionID uint64, epoch *metapb.RegionEpoch) raft_cmdpb.RaftCmdRequest {
	req := raft_cmdpb.RaftCmdRequest{}
	req.Header = &raft_cmdpb.RaftRequestHeader{RegionId: regionID, RegionEpoch: epoch}
	return req
}

func NewRequest(regionID uint64, epoch *metapb.RegionEpoch, requests []*raft_cmdpb.Request) raft_cmdpb.RaftCmdRequest {
	req := NewBaseRequest(regionID, epoch)
	req.Requests = requests
	return req
}

func NewAdminRequest(regionID uint64, epoch *metapb.RegionEpoch, request *raft_cmdpb.AdminRequest) *raft_cmdpb.RaftCmdRequest {
	req := NewBaseRequest(regionID, epoch)
	req.AdminRequest = request
	return &req
}

func NewPutCfCmd(cf string, key, value []byte) *raft_cmdpb.Request {
	cmd := &raft_cmdpb.Request{}
	cmd.CmdType = raft_cmdpb.CmdType_Put
	cmd.Put = &raft_cmdpb.PutRequest{Key: key, Value: value, Cf: cf}
	return cmd
}

func NewGetCfCmd(cf string, key []byte) *raft_cmdpb.Request {
	get := &raft_cmdpb.GetRequest{
		Cf:  cf,
		Key: key,
	}
	cmd := &raft_cmdpb.Request{
		CmdType: raft_cmdpb.CmdType_Get,
		Get:     get,
	}
	return cmd
}

func NewDeleteCfCmd(cf string, key []byte) *raft_cmdpb.Request {
	delete := &raft_cmdpb.DeleteRequest{
		Cf:  cf,
		Key: key,
	}
	cmd := &raft_cmdpb.Request{
		CmdType: raft_cmdpb.CmdType_Delete,
		Delete:  delete,
	}
	return cmd
}

func NewSnapCmd() *raft_cmdpb.Request {
	cmd := &raft_cmdpb.Request{
		CmdType: raft_cmdpb.CmdType_Snap,
		Snap:    &raft_cmdpb.SnapRequest{},
	}
	return cmd
}

func NewTransferLeaderCmd(peer *metapb.Peer) *raft_cmdpb.AdminRequest {
	transferLeader := raft_cmdpb.TransferLeaderRequest{Peer: peer}
	cmd := &raft_cmdpb.AdminRequest{
		CmdType:        raft_cmdpb.AdminCmdType_TransferLeader,
		TransferLeader: &transferLeader,
	}
	return cmd
}

func MustGetCf(engine *engine_util.Engines, cf string, key []byte, value []byte) {
	for i := 0; i < 300; i++ {
		val, err := engine_util.GetCF(engine.Kv, cf, key)
		if err == nil && (value == nil || bytes.Compare(val, value) == 0) {
			return
		}
		SleepMS(20)
	}
	panic(fmt.Sprintf("can't get value %s for key %s", hex.EncodeToString(value), hex.EncodeToString(key)))
}

func MustGetCfEqual(engine *engine_util.Engines, cf string, key []byte, value []byte) {
	MustGetCf(engine, cf, key, value)
}

func MustGetEqual(engine *engine_util.Engines, key []byte, value []byte) {
	MustGetCf(engine, engine_util.CfDefault, key, value)
}

func MustGetCfNone(engine *engine_util.Engines, cf string, key []byte) {
	var val []byte
	var err error
	for i := 0; i < 300; i++ {
		val, err = engine_util.GetCF(engine.Kv, cf, key)
		if err == badger.ErrKeyNotFound {
			return
		}
		SleepMS(20)
	}
	panic(fmt.Sprintf("get value %s for key %s", hex.EncodeToString(val), hex.EncodeToString(key)))
}

func MustGetNone(engine *engine_util.Engines, key []byte) {
	MustGetCfNone(engine, engine_util.CfDefault, key)
}

func NewTestCluster(count int, cfg *config.Config) *Cluster {
	log.SetLevelByString(cfg.LogLevel)
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds | log.Lshortfile)
	schedulerClient := NewMockSchedulerClient(0, uint64(count)+1)
	simulator := NewNodeSimulator(schedulerClient)
	return NewCluster(count, schedulerClient, simulator, cfg)
}