move core to pkg

Signed-off-by: Ryan Leung <rleungx@gmail.com>

pkg/autoscaling/calculation.go

@@ -23,11 +23,11 @@ import (
 	"github.com/pingcap/errors"
 	"github.com/pingcap/log"
 	promClient "github.com/prometheus/client_golang/api"
+	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/errs"
 	"github.com/tikv/pd/pkg/utils/typeutil"
 	"github.com/tikv/pd/server/cluster"
 	"github.com/tikv/pd/server/config"
-	"github.com/tikv/pd/server/core"
 	"github.com/tikv/pd/server/schedule/filter"
 	"go.etcd.io/etcd/clientv3"
 	"go.uber.org/zap"

pkg/autoscaling/calculation_test.go

@@ -23,9 +23,9 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/mock/mockcluster"
 	"github.com/tikv/pd/server/config"
-	"github.com/tikv/pd/server/core"
 )
 
 func TestGetScaledTiKVGroups(t *testing.T) {

server/core/basic_cluster.go → pkg/core/basic_cluster.go

@@ -16,8 +16,8 @@ package core
 
 import (
 	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/tikv/pd/pkg/core/storelimit"
 	"github.com/tikv/pd/pkg/utils/syncutil"
-	"github.com/tikv/pd/server/core/storelimit"
 )
 
 // BasicCluster provides basic data member and interface for a tikv cluster.

server/core/region.go → pkg/core/region.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"encoding/hex"
 	"fmt"
+	"math"
 	"reflect"
 	"sort"
 	"strings"
@@ -1674,3 +1675,71 @@ func NeedTransferWitnessLeader(region *RegionInfo) bool {
 	}
 	return region.GetLeader().IsWitness
 }
+
+// SplitRegions split a set of RegionInfo by the middle of regionKey. Only for test purpose.
+func SplitRegions(regions []*RegionInfo) []*RegionInfo {
+	results := make([]*RegionInfo, 0, len(regions)*2)
+	for _, region := range regions {
+		start, end := byte(0), byte(math.MaxUint8)
+		if len(region.GetStartKey()) > 0 {
+			start = region.GetStartKey()[0]
+		}
+		if len(region.GetEndKey()) > 0 {
+			end = region.GetEndKey()[0]
+		}
+		middle := []byte{start/2 + end/2}
+		left := region.Clone()
+		left.meta.Id = region.GetID() + uint64(len(regions))
+		left.meta.EndKey = middle
+		left.meta.RegionEpoch.Version++
+		right := region.Clone()
+		right.meta.Id = region.GetID() + uint64(len(regions)*2)
+		right.meta.StartKey = middle
+		right.meta.RegionEpoch.Version++
+		results = append(results, left, right)
+	}
+	return results
+}
+
+// MergeRegions merge a set of RegionInfo by regionKey. Only for test purpose.
+func MergeRegions(regions []*RegionInfo) []*RegionInfo {
+	results := make([]*RegionInfo, 0, len(regions)/2)
+	for i := 0; i < len(regions); i += 2 {
+		left := regions[i]
+		right := regions[i]
+		if i+1 < len(regions) {
+			right = regions[i+1]
+		}
+		region := &RegionInfo{meta: &metapb.Region{
+			Id:       left.GetID(),
+			StartKey: left.GetStartKey(),
+			EndKey:   right.GetEndKey(),
+			Peers:    left.meta.Peers,
+		}}
+		if left.GetRegionEpoch().GetVersion() > right.GetRegionEpoch().GetVersion() {
+			region.meta.RegionEpoch = left.GetRegionEpoch()
+		} else {
+			region.meta.RegionEpoch = right.GetRegionEpoch()
+		}
+		region.meta.RegionEpoch.Version++
+		region.leader = left.leader
+		results = append(results, region)
+	}
+	return results
+}
+
+// NewTestRegionInfo creates a new RegionInfo for test purpose.
+func NewTestRegionInfo(regionID, storeID uint64, start, end []byte, opts ...RegionCreateOption) *RegionInfo {
+	leader := &metapb.Peer{
+		Id:      regionID,
+		StoreId: storeID,
+	}
+	metaRegion := &metapb.Region{
+		Id:          regionID,
+		StartKey:    start,
+		EndKey:      end,
+		Peers:       []*metapb.Peer{leader},
+		RegionEpoch: &metapb.RegionEpoch{ConfVer: 1, Version: 1},
+	}
+	return NewRegionInfo(metaRegion, leader, opts...)
+}

server/core/region_option.go → pkg/core/region_option.go

@@ -16,7 +16,9 @@ package core
 
 import (
 	"math"
+	"math/rand"
 	"sort"
+	"time"
 
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
@@ -237,15 +239,6 @@ func SetWrittenQuery(v uint64) RegionCreateOption {
 	return SetQueryStats(q)
 }
 
-// SetQueryNum sets the read and write query with specific num.
-// This func is only used for test and simulator.
-func SetQueryNum(read, write uint64) RegionCreateOption {
-	r := RandomKindReadQuery(read)
-	w := RandomKindWriteQuery(write)
-	q := mergeQueryStat(r, w)
-	return SetQueryStats(q)
-}
-
 // SetQueryStats sets the query stats for the region, it will cover previous statistic.
 // This func is only used for unit test.
 func SetQueryStats(v *pdpb.QueryStats) RegionCreateOption {
@@ -369,6 +362,64 @@ func SetFromHeartbeat(fromHeartbeat bool) RegionCreateOption {
 	}
 }
 
+// RandomKindReadQuery returns query stat with random query kind, only used for unit test.
+func RandomKindReadQuery(queryRead uint64) *pdpb.QueryStats {
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	switch r.Intn(3) {
+	case 0:
+		return &pdpb.QueryStats{
+			Coprocessor: queryRead,
+		}
+	case 1:
+		return &pdpb.QueryStats{
+			Scan: queryRead,
+		}
+	case 2:
+		return &pdpb.QueryStats{
+			Get: queryRead,
+		}
+	default:
+		return &pdpb.QueryStats{}
+	}
+}
+
+// RandomKindWriteQuery returns query stat with random query kind, only used for unit test.
+func RandomKindWriteQuery(queryWrite uint64) *pdpb.QueryStats {
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	switch r.Intn(7) {
+	case 0:
+		return &pdpb.QueryStats{
+			Put: queryWrite,
+		}
+	case 1:
+		return &pdpb.QueryStats{
+			Delete: queryWrite,
+		}
+	case 2:
+		return &pdpb.QueryStats{
+			DeleteRange: queryWrite,
+		}
+	case 3:
+		return &pdpb.QueryStats{
+			AcquirePessimisticLock: queryWrite,
+		}
+	case 4:
+		return &pdpb.QueryStats{
+			Rollback: queryWrite,
+		}
+	case 5:
+		return &pdpb.QueryStats{
+			Prewrite: queryWrite,
+		}
+	case 6:
+		return &pdpb.QueryStats{
+			Commit: queryWrite,
+		}
+	default:
+		return &pdpb.QueryStats{}
+	}
+}
+
 func mergeQueryStat(q1, q2 *pdpb.QueryStats) *pdpb.QueryStats {
 	if q1 == nil && q2 == nil {
 		return &pdpb.QueryStats{}

server/core/region_test.go → pkg/core/region_test.go

@@ -615,7 +615,7 @@ func checkRegions(re *require.Assertions, regions *RegionsInfo) {
 }
 
 func BenchmarkUpdateBuckets(b *testing.B) {
-	region := NewTestRegionInfo([]byte{}, []byte{})
+	region := NewTestRegionInfo(1, 1, []byte{}, []byte{})
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		buckets := &metapb.Buckets{RegionId: 0, Version: uint64(i)}

server/core/region_tree_test.go → pkg/core/region_tree_test.go

@@ -114,7 +114,7 @@ func TestRegionItem(t *testing.T) {
 }
 
 func newRegionWithStat(start, end string, size, keys int64) *RegionInfo {
-	region := NewTestRegionInfo([]byte(start), []byte(end))
+	region := NewTestRegionInfo(1, 1, []byte(start), []byte(end))
 	region.approximateSize, region.approximateKeys = size, keys
 	return region
 }
@@ -151,10 +151,10 @@ func TestRegionTree(t *testing.T) {
 
 	re.Nil(tree.search([]byte("a")))
 
-	regionA := NewTestRegionInfo([]byte("a"), []byte("b"))
-	regionB := NewTestRegionInfo([]byte("b"), []byte("c"))
-	regionC := NewTestRegionInfo([]byte("c"), []byte("d"))
-	regionD := NewTestRegionInfo([]byte("d"), []byte{})
+	regionA := NewTestRegionInfo(1, 1, []byte("a"), []byte("b"))
+	regionB := NewTestRegionInfo(2, 2, []byte("b"), []byte("c"))
+	regionC := NewTestRegionInfo(3, 3, []byte("c"), []byte("d"))
+	regionD := NewTestRegionInfo(4, 4, []byte("d"), []byte{})
 
 	updateNewItem(tree, regionA)
 	updateNewItem(tree, regionC)
@@ -274,14 +274,14 @@ func TestRandomRegion(t *testing.T) {
 	r := tree.RandomRegion(nil)
 	re.Nil(r)
 
-	regionA := NewTestRegionInfo([]byte(""), []byte("g"))
+	regionA := NewTestRegionInfo(1, 1, []byte(""), []byte("g"))
 	updateNewItem(tree, regionA)
 	ra := tree.RandomRegion([]KeyRange{NewKeyRange("", "")})
 	re.Equal(regionA, ra)
 
-	regionB := NewTestRegionInfo([]byte("g"), []byte("n"))
-	regionC := NewTestRegionInfo([]byte("n"), []byte("t"))
-	regionD := NewTestRegionInfo([]byte("t"), []byte(""))
+	regionB := NewTestRegionInfo(2, 2, []byte("g"), []byte("n"))
+	regionC := NewTestRegionInfo(3, 3, []byte("n"), []byte("t"))
+	regionD := NewTestRegionInfo(4, 4, []byte("t"), []byte(""))
 	updateNewItem(tree, regionB)
 	updateNewItem(tree, regionC)
 	updateNewItem(tree, regionD)
@@ -316,7 +316,7 @@ func TestRandomRegionDiscontinuous(t *testing.T) {
 	re.Nil(r)
 
 	// test for single region
-	regionA := NewTestRegionInfo([]byte("c"), []byte("f"))
+	regionA := NewTestRegionInfo(1, 1, []byte("c"), []byte("f"))
 	updateNewItem(tree, regionA)
 	ra := tree.RandomRegion([]KeyRange{NewKeyRange("c", "e")})
 	re.Nil(ra)
@@ -331,7 +331,7 @@ func TestRandomRegionDiscontinuous(t *testing.T) {
 	ra = tree.RandomRegion([]KeyRange{NewKeyRange("a", "g")})
 	re.Equal(regionA, ra)
 
-	regionB := NewTestRegionInfo([]byte("n"), []byte("x"))
+	regionB := NewTestRegionInfo(2, 2, []byte("n"), []byte("x"))
 	updateNewItem(tree, regionB)
 	rb := tree.RandomRegion([]KeyRange{NewKeyRange("g", "x")})
 	re.Equal(regionB, rb)
@@ -342,11 +342,11 @@ func TestRandomRegionDiscontinuous(t *testing.T) {
 	rb = tree.RandomRegion([]KeyRange{NewKeyRange("o", "y")})
 	re.Nil(rb)
 
-	regionC := NewTestRegionInfo([]byte("z"), []byte(""))
+	regionC := NewTestRegionInfo(3, 3, []byte("z"), []byte(""))
 	updateNewItem(tree, regionC)
 	rc := tree.RandomRegion([]KeyRange{NewKeyRange("y", "")})
 	re.Equal(regionC, rc)
-	regionD := NewTestRegionInfo([]byte(""), []byte("a"))
+	regionD := NewTestRegionInfo(4, 4, []byte(""), []byte("a"))
 	updateNewItem(tree, regionD)
 	rd := tree.RandomRegion([]KeyRange{NewKeyRange("", "b")})
 	re.Equal(regionD, rd)
@@ -379,7 +379,7 @@ func checkRandomRegion(re *require.Assertions, tree *regionTree, regions []*Regi
 }
 
 func newRegionItem(start, end []byte) *regionItem {
-	return &regionItem{RegionInfo: NewTestRegionInfo(start, end)}
+	return &regionItem{RegionInfo: NewTestRegionInfo(1, 1, start, end)}
 }
 
 type mockRegionTreeData struct {

server/core/store.go → pkg/core/store.go

@@ -22,9 +22,9 @@ import (
 	"github.com/docker/go-units"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/log"
+	"github.com/tikv/pd/pkg/core/storelimit"
 	"github.com/tikv/pd/pkg/errs"
 	"github.com/tikv/pd/pkg/utils/typeutil"
-	"github.com/tikv/pd/server/core/storelimit"
 	"go.uber.org/zap"
 )
 
@@ -80,6 +80,24 @@ func NewStoreInfo(store *metapb.Store, opts ...StoreCreateOption) *StoreInfo {
 	return storeInfo
 }
 
+// NewStoreInfoWithLabel is create a store with specified labels. only for test purpose.
+func NewStoreInfoWithLabel(id uint64, labels map[string]string) *StoreInfo {
+	storeLabels := make([]*metapb.StoreLabel, 0, len(labels))
+	for k, v := range labels {
+		storeLabels = append(storeLabels, &metapb.StoreLabel{
+			Key:   k,
+			Value: v,
+		})
+	}
+	store := NewStoreInfo(
+		&metapb.Store{
+			Id:     id,
+			Labels: storeLabels,
+		},
+	)
+	return store
+}
+
 // Clone creates a copy of current StoreInfo.
 func (s *StoreInfo) Clone(opts ...StoreCreateOption) *StoreInfo {
 	store := *s

server/core/store_option.go → pkg/core/store_option.go

@@ -19,8 +19,8 @@ import (
 
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
+	"github.com/tikv/pd/pkg/core/storelimit"
 	"github.com/tikv/pd/pkg/utils/typeutil"
-	"github.com/tikv/pd/server/core/storelimit"
 )
 
 // StoreCreateOption is used to create store.