cli: add an option to set node localities in cockroach demo

pkg/cli/cliflags/flags.go

@@ -889,6 +889,21 @@ The line length where sqlfmt will try to wrap.`,
 		Description: `How many in-memory nodes to create for the demo.`,
 	}
 
+	DemoNodeLocality = FlagInfo{
+		Name: "demo-locality",
+		Description: `
+Locality information for each demo node. The input is a comma separated
+list of key-value pairs, where the i'th pair is the locality setting
+for the i'th demo cockroach node. For example:
+<PRE>
+
+	--demo-locality=region=us-east1,region=us-east2,region=us-east3
+
+Assigns node 1's region to us-east1, node 2's region to us-east2,
+and node 3's region to us-east3.
+`,
+	}
+
 	UseEmptyDatabase = FlagInfo{
 		Name: "empty",
 		Description: `

pkg/cli/context.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/config"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -145,6 +146,7 @@ func initCLIDefaults() {
 
 	demoCtx.nodes = 1
 	demoCtx.useEmptyDatabase = false
+	demoCtx.localities = roachpb.Locality{}
 
 	initPreFlagsDefaults()
 
@@ -334,4 +336,5 @@ var sqlfmtCtx struct {
 var demoCtx struct {
 	nodes            int
 	useEmptyDatabase bool
+	localities       roachpb.Locality
 }

pkg/cli/demo.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/cli/cliflags"
 	"github.com/cockroachdb/cockroach/pkg/config"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
@@ -28,6 +29,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/workload"
 	"github.com/cockroachdb/cockroach/pkg/workload/workloadsql"
 	"github.com/gogo/protobuf/proto"
+	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 )
@@ -87,6 +89,20 @@ func setupTransientServers(
 	cleanup = func() {}
 	ctx := context.Background()
 
+	if demoCtx.nodes <= 0 {
+		return "", "", cleanup, errors.Errorf("must have a positive number of nodes")
+	}
+
+	var nodeLocalityTiers []roachpb.Tier
+	if len(demoCtx.localities.Tiers) != 0 {
+		// Error out of localities don't line up with requested node
+		// count before doing any sort of setup.
+		if len(demoCtx.localities.Tiers) != demoCtx.nodes {
+			return "", "", cleanup, errors.Errorf("number of localities specified must equal number of nodes")
+		}
+		nodeLocalityTiers = demoCtx.localities.Tiers
+	}
+
 	// Set up logging. For demo/transient server we use non-standard
 	// behavior where we avoid file creation if possible.
 	df := cmd.Flags().Lookup(cliflags.LogDir.Name)
@@ -129,17 +145,25 @@ func setupTransientServers(
 		},
 		Stopper: stopper,
 	}
+
 	serverFactory := server.TestServerFactory
-	s := serverFactory.New(args).(*server.TestServer)
-	if err := s.Start(args); err != nil {
-		return connURL, adminURL, cleanup, err
-	}
-	args.JoinAddr = s.ServingRPCAddr()
-	for i := 0; i < demoCtx.nodes-1; i++ {
-		s := serverFactory.New(args).(*server.TestServer)
-		if err := s.Start(args); err != nil {
+	var s *server.TestServer
+	for i := 0; i < demoCtx.nodes; i++ {
+		// All the nodes connect to the address of the first server created.
+		if s != nil {
+			args.JoinAddr = s.ServingRPCAddr()
+		}
+		if nodeLocalityTiers != nil {
+			args.Locality = roachpb.Locality{Tiers: nodeLocalityTiers[i : i+1]}
+		}
+		serv := serverFactory.New(args).(*server.TestServer)
+		if err := serv.Start(args); err != nil {
 			return connURL, adminURL, cleanup, err
 		}
+		// Remember the first server created.
+		if i == 0 {
+			s = serv
+		}
 	}
 
 	// Prepare the URL for use by the SQL shell.

pkg/cli/flags.go

@@ -592,6 +592,7 @@ func init() {
 	// The --empty flag is only valid for the top level demo command,
 	// so we use the regular flag set.
 	BoolFlag(demoCmd.Flags(), &demoCtx.useEmptyDatabase, cliflags.UseEmptyDatabase, false)
+	VarFlag(demoFlags, &demoCtx.localities, cliflags.DemoNodeLocality)
 
 	// sqlfmt command.
 	fmtFlags := sqlfmtCmd.Flags()