Merge pull request #12516 from BramGruneir/allocsim

cmd: add the ability to specify localities in allocsim

pkg/cmd/allocsim/main.go

@@ -33,6 +33,7 @@ import (
 	"golang.org/x/net/context"
 
 	"github.com/cockroachdb/cockroach/pkg/cmd/internal/localcluster"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -43,6 +44,7 @@ var workers = flag.Int("w", 1, "number of workers; the i'th worker talks to node
 var numNodes = flag.Int("n", 4, "number of nodes")
 var duration = flag.Duration("duration", math.MaxInt64, "how long to run the simulation for")
 var blockSize = flag.Int("b", 1000, "block size")
+var numLocalities = flag.Int("l", 0, "number of localities")
 
 func newRand() *rand.Rand {
 	return rand.New(rand.NewSource(timeutil.Now().UnixNano()))
@@ -53,8 +55,7 @@ func newRand() *rand.Rand {
 // talks to node i%numNodes. Every second a monitor goroutine outputs status
 // such as the per-node replica and leaseholder counts.
 //
-// TODO(peter): Allow configuration of per-node locality settings and
-// zone-config constraints.
+// TODO(peter): Allow configuration zone-config constraints.
 
 type allocSim struct {
 	*localcluster.Cluster
@@ -68,6 +69,7 @@ type allocSim struct {
 		replicas []int
 		leases   []int
 	}
+	localities []roachpb.Locality
 }
 
 func newAllocSim(c *localcluster.Cluster) *allocSim {
@@ -183,11 +185,14 @@ func (a *allocSim) rangeStats(d time.Duration) {
 
 const padding = "__________"
 
-func formatHeader(header string, numberNodes int) string {
+func formatHeader(header string, numberNodes int, localities []roachpb.Locality) string {
 	var buf bytes.Buffer
 	_, _ = buf.WriteString(header)
 	for i := 1; i <= numberNodes; i++ {
 		node := fmt.Sprintf("%d", i)
+		if localities != nil {
+			node += fmt.Sprintf(":%s", localities[i-1])
+		}
 		fmt.Fprintf(&buf, "%s%s", padding[:len(padding)-len(node)], node)
 	}
 	return buf.String()
@@ -229,7 +234,7 @@ func (a *allocSim) monitor(d time.Duration) {
 
 		if ticks%20 == 0 || numReplicas != len(replicas) {
 			numReplicas = len(replicas)
-			fmt.Println(formatHeader("_elapsed__ops/sec___errors_replicas", numReplicas))
+			fmt.Println(formatHeader("_elapsed__ops/sec___errors_replicas", numReplicas, a.localities))
 		}
 
 		fmt.Printf("%8s %8.1f %8d %8d%s\n",
@@ -248,7 +253,7 @@ func (a *allocSim) finalStatus() {
 	// TODO(bram): With the addition of localities, these stats will have to be
 	// updated.
 
-	fmt.Println(formatHeader("___stats___________________________", len(a.ranges.replicas)))
+	fmt.Println(formatHeader("___stats___________________________", len(a.ranges.replicas), a.localities))
 
 	genStats := func(name string, counts []int) {
 		var total float64
@@ -287,6 +292,24 @@ func main() {
 	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
 	a := newAllocSim(c)
 
+	var perNodeArgs map[int][]string
+	if *numLocalities > 0 {
+		perNodeArgs = make(map[int][]string)
+		a.localities = make([]roachpb.Locality, len(c.Nodes), len(c.Nodes))
+		for i := range c.Nodes {
+			locality := roachpb.Locality{
+				Tiers: []roachpb.Tier{
+					{
+						Key:   "l",
+						Value: fmt.Sprintf("%d", i%*numLocalities),
+					},
+				},
+			}
+			perNodeArgs[i] = []string{fmt.Sprintf("--locality=%s", locality)}
+			a.localities[i] = locality
+		}
+	}
+
 	go func() {
 		var exitStatus int
 		select {
@@ -296,12 +319,12 @@ func main() {
 		case <-time.After(*duration):
 			log.Infof(context.Background(), "finished run of: %s", *duration)
 		}
-		a.finalStatus()
 		c.Close()
+		a.finalStatus()
 		os.Exit(exitStatus)
 	}()
 
-	c.Start("allocsim", *workers, flag.Args(), []string{})
+	c.Start("allocsim", *workers, nil, flag.Args(), perNodeArgs)
 	c.UpdateZoneConfig(1, 1<<20)
 	a.run(*workers)
 }

pkg/cmd/internal/localcluster/localcluster.go

@@ -30,6 +30,11 @@ import (
 	"text/tabwriter"
 	"time"
 
+	"github.com/pkg/errors"
+	// Import postgres driver.
+	_ "github.com/lib/pq"
+	"golang.org/x/net/context"
+
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/config"
 	"github.com/cockroachdb/cockroach/pkg/internal/client"
@@ -44,11 +49,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
-
-	"github.com/pkg/errors"
-	// Import postgres driver.
-	_ "github.com/lib/pq"
-	"golang.org/x/net/context"
 )
 
 const basePort = 26257
@@ -92,9 +92,13 @@ func New(size int) *Cluster {
 }
 
 // Start starts a cluster. The numWorkers parameter controls the SQL connection
-// settings to avoid unnecessary connection creation. The args parameter can be
-// used to pass extra arguments to each node.
-func (c *Cluster) Start(db string, numWorkers int, args, env []string) {
+// settings to avoid unnecessary connection creation. The allNodeArgs parameter
+// can be used to pass extra arguments to every node. The perNodeArgs parameter
+// can be used to pass extra arguments to an individual node. If not nil, its
+// size must equal the number of nodes.
+func (c *Cluster) Start(
+	db string, numWorkers int, env, allNodeArgs []string, perNodeArgs map[int][]string,
+) {
 	c.started = timeutil.Now()
 
 	baseCtx := &base.Config{
@@ -104,8 +108,13 @@ func (c *Cluster) Start(db string, numWorkers int, args, env []string) {
 	c.rpcCtx = rpc.NewContext(log.AmbientContext{}, baseCtx,
 		hlc.NewClock(hlc.UnixNano, 0), c.stopper)
 
+	if perNodeArgs != nil && len(perNodeArgs) != len(c.Nodes) {
+		panic(fmt.Sprintf("there are %d nodes, but perNodeArgs' length is %d",
+			len(c.Nodes), len(perNodeArgs)))
+	}
+
 	for i := range c.Nodes {
-		c.Nodes[i] = c.makeNode(i, args, env)
+		c.Nodes[i] = c.makeNode(i, append(append([]string(nil), allNodeArgs...), perNodeArgs[i]...), env)
 		c.Clients[i] = c.makeClient(i)
 		c.Status[i] = c.makeStatus(i)
 		c.DB[i] = c.makeDB(i, numWorkers, db)

pkg/cmd/zerosum/main.go

@@ -454,7 +454,7 @@ func main() {
 		os.Exit(1)
 	}()
 
-	c.Start("zerosum", *workers, flag.Args(), nil)
+	c.Start("zerosum", *workers, nil, flag.Args(), nil)
 
 	z := newZeroSum(c, *numAccounts, *chaosType)
 	z.run(*workers, *monkeys)