diff --git a/pkg/cli/cliflags/flags.go b/pkg/cli/cliflags/flags.go index ea846cb28186..639a48bf7619 100644 --- a/pkg/cli/cliflags/flags.go +++ b/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: +
+ + --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: ` diff --git a/pkg/cli/context.go b/pkg/cli/context.go index 3f7c9263302e..8c812e6b0150 100644 --- a/pkg/cli/context.go +++ b/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 } diff --git a/pkg/cli/demo.go b/pkg/cli/demo.go index 6af08b0f4394..b6c2ea7e483b 100644 --- a/pkg/cli/demo.go +++ b/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. diff --git a/pkg/cli/flags.go b/pkg/cli/flags.go index 04d32074f62b..90191b7d64cf 100644 --- a/pkg/cli/flags.go +++ b/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()