simulator: support for more complex configurations.

cmd/simulator/main.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"flag"
 	"os"
 	"os/signal"
 	"syscall"
@@ -30,13 +31,17 @@ import (
 	_ "github.com/pingcap/pd/table"
 )
 
+var confName = flag.String("conf", "", "config name")
+
 func main() {
+	flag.Parse()
+
 	_, local, clean := NewSingleServer()
 	err := local.Run()
 	if err != nil {
 		log.Fatal("run server error:", err)
 	}
-	driver := faketikv.NewDriver(local.GetAddr())
+	driver := faketikv.NewDriver(local.GetAddr(), *confName)
 	err = driver.Prepare()
 	if err != nil {
 		log.Fatal("simulator prepare error:", err)
@@ -49,16 +54,26 @@ func main() {
 		syscall.SIGTERM,
 		syscall.SIGQUIT)
 
+	simResult := "FAIL"
+
+EXIT:
 	for {
 		select {
 		case <-tick.C:
 			driver.Tick()
+			if driver.Check() {
+				simResult = "OK"
+				break EXIT
+			}
 		case <-sc:
-			driver.Stop()
-			clean()
-			return
+			break EXIT
 		}
 	}
+
+	driver.Stop()
+	clean()
+
+	log.Infof("Simulation finish. Conf: %s, TotalTick: %d, Result: %s", *confName, driver.TickCount(), simResult)
 }
 
 // NewSingleServer creates a pd server for simulator.

pkg/faketikv/cases/balance_leader.go

@@ -0,0 +1,62 @@
+// Copyright 2017 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 cases
+
+import (
+	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/pingcap/pd/server/core"
+	log "github.com/sirupsen/logrus"
+)
+
+func newBalanceLeader() *Conf {
+	var conf Conf
+
+	for i := 1; i <= 3; i++ {
+		conf.Stores = append(conf.Stores, Store{
+			ID:        uint64(i),
+			Status:    metapb.StoreState_Up,
+			Capacity:  10 * gb,
+			Available: 9 * gb,
+		})
+	}
+
+	var id idAllocator
+	id.setMaxID(3)
+	for i := 0; i < 1000; i++ {
+		peers := []*metapb.Peer{
+			{Id: id.nextID(), StoreId: 1},
+			{Id: id.nextID(), StoreId: 2},
+			{Id: id.nextID(), StoreId: 3},
+		}
+		conf.Regions = append(conf.Regions, Region{
+			ID:     id.nextID(),
+			Peers:  peers,
+			Leader: peers[0],
+			Size:   96 * mb,
+		})
+	}
+	conf.MaxID = id.maxID
+
+	conf.Checker = func(regions *core.RegionsInfo) bool {
+		count1 := regions.GetStoreLeaderCount(1)
+		count2 := regions.GetStoreLeaderCount(2)
+		count3 := regions.GetStoreLeaderCount(3)
+		log.Infof("leader counts: %v %v %v", count1, count2, count3)
+
+		return count1 <= 350 &&
+			count2 >= 300 &&
+			count3 >= 300
+	}
+	return &conf
+}

pkg/faketikv/cases/cases.go

@@ -0,0 +1,78 @@
+// Copyright 2017 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 cases
+
+import (
+	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/pingcap/pd/server/core"
+)
+
+// Store is the config to simulate tikv.
+type Store struct {
+	ID           uint64
+	Status       metapb.StoreState
+	Labels       []metapb.StoreLabel
+	Capacity     uint64
+	Available    uint64
+	LeaderWeight float32
+	RegionWeight float32
+}
+
+// Region is the config to simulate a region.
+type Region struct {
+	ID     uint64
+	Peers  []*metapb.Peer
+	Leader *metapb.Peer
+	Size   int64
+}
+
+// Conf represents a test suite for simulator.
+type Conf struct {
+	Stores  []Store
+	Regions []Region
+	MaxID   uint64
+
+	Checker func(*core.RegionsInfo) bool // To check the schedule is finished.
+}
+
+const (
+	kb = 1024
+	mb = kb * 1024
+	gb = mb * 1024
+)
+
+type idAllocator struct {
+	maxID uint64
+}
+
+func (a *idAllocator) nextID() uint64 {
+	a.maxID++
+	return a.maxID
+}
+
+func (a *idAllocator) setMaxID(id uint64) {
+	a.maxID = id
+}
+
+var confMap = map[string]func() *Conf{
+	"balance-leader": newBalanceLeader,
+}
+
+// NewConf creates a config to initialize simulator cluster.
+func NewConf(name string) *Conf {
+	if f, ok := confMap[name]; ok {
+		return f()
+	}
+	return nil
+}

pkg/faketikv/cluster.go

@@ -14,23 +14,77 @@
 package faketikv
 
 import (
+	"fmt"
+	"math/rand"
+	"sort"
+
+	"github.com/juju/errors"
 	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/pingcap/pd/pkg/faketikv/cases"
 	"github.com/pingcap/pd/server/core"
 	log "github.com/sirupsen/logrus"
 )
 
 // ClusterInfo records all cluster information.
 type ClusterInfo struct {
 	*core.RegionsInfo
-	Nodes       map[uint64]*Node
-	firstRegion *core.RegionInfo
+	Nodes map[uint64]*Node
+}
+
+// NewClusterInfo creates the initialized cluster with config.
+func NewClusterInfo(pdAddr string, conf *cases.Conf) (*ClusterInfo, error) {
+	cluster := &ClusterInfo{
+		RegionsInfo: core.NewRegionsInfo(),
+		Nodes:       make(map[uint64]*Node),
+	}
+
+	for _, store := range conf.Stores {
+		node, err := NewNode(store.ID, fmt.Sprintf("mock:://tikv-%d", store.ID), pdAddr)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		node.clusterInfo = cluster
+		cluster.Nodes[store.ID] = node
+	}
+
+	splitKeys := generateKeys(len(conf.Regions) - 1)
+	for i, region := range conf.Regions {
+		meta := &metapb.Region{
+			Id:          region.ID,
+			Peers:       region.Peers,
+			RegionEpoch: &metapb.RegionEpoch{ConfVer: 1, Version: 1},
+		}
+		if i > 0 {
+			meta.StartKey = []byte(splitKeys[i-1])
+		}
+		if i < len(conf.Regions)-1 {
+			meta.EndKey = []byte(splitKeys[i])
+		}
+		regionInfo := core.NewRegionInfo(meta, region.Leader)
+		regionInfo.ApproximateSize = region.Size
+		cluster.RegionsInfo.SetRegion(regionInfo)
+	}
+
+	return cluster, nil
 }
 
-// GetBootstrapInfo returns first region and its leader store.
-func (c *ClusterInfo) GetBootstrapInfo() (*metapb.Store, *metapb.Region) {
-	storeID := c.firstRegion.Leader.GetStoreId()
-	store := c.Nodes[storeID]
-	return store.Store, c.firstRegion.Region
+// GetBootstrapInfo returns a valid bootstrap store and region.
+func (c *ClusterInfo) GetBootstrapInfo() (*metapb.Store, *metapb.Region, error) {
+	region := c.RegionsInfo.RandRegion()
+	if region == nil {
+		return nil, nil, errors.New("no region found for bootstrap")
+	}
+	if region.Leader == nil {
+		return nil, nil, errors.New("bootstrap region has no leader")
+	}
+	store := c.Nodes[region.Leader.GetStoreId()]
+	if store == nil {
+		return nil, nil, errors.Errorf("bootstrap store %v not found", region.Leader.GetStoreId())
+	}
+	region.StartKey, region.EndKey = []byte(""), []byte("")
+	region.RegionEpoch = &metapb.RegionEpoch{}
+	region.Peers = []*metapb.Peer{region.Leader}
+	return store.Store, region.Region, nil
 }
 
 func (c *ClusterInfo) nodeHealth(storeID uint64) bool {
@@ -42,7 +96,7 @@ func (c *ClusterInfo) nodeHealth(storeID uint64) bool {
 	return n.GetState() == Up
 }
 
-func (c ClusterInfo) electNewLeader(region *core.RegionInfo) *metapb.Peer {
+func (c *ClusterInfo) electNewLeader(region *core.RegionInfo) *metapb.Peer {
 	var (
 		unhealth         int
 		newLeaderStoreID uint64
@@ -103,3 +157,28 @@ func (c *ClusterInfo) AddTask(task Task) {
 		n.AddTask(task)
 	}
 }
+
+const (
+	// 26^10 ~= 1.4e+14, should be enough.
+	keyChars = "abcdefghijklmnopqrstuvwxyz"
+	keyLen   = 10
+)
+
+// generate ordered, unique strings.
+func generateKeys(size int) []string {
+	m := make(map[string]struct{}, size)
+	for len(m) < size {
+		k := make([]byte, keyLen)
+		for i := range k {
+			k[i] = keyChars[rand.Intn(len(keyChars))]
+		}
+		m[string(k)] = struct{}{}
+	}
+
+	v := make([]string, 0, size)
+	for k := range m {
+		v = append(v, k)
+	}
+	sort.Sort(sort.StringSlice(v))
+	return v
+}