Merge #39786

39786: cli: Extending locality assignment for cockroach demo nodes r=rohany a=rohany

This PR allows for setting multiple tiers of locality
for cockroach demo nodes, as opposed to previous work
allowing for only a single tier to be set for each node.

Now, use a colon to separate the locality information
for each demo node. For example,

```
--demo-locality=region=us-east1,az=1:region=us-east1,az=2:region=us-east1,az=3
```

Fixes #39780.

Release note (cli change): Extends the cockroach demo node locality setting to allow for multiple locality tiers per demo node.

Co-authored-by: Rohan Yadav <[email protected]>

pkg/cli/cliflags/flags.go

@@ -900,15 +900,16 @@ The line length where sqlfmt will try to wrap.`,
 	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:
+Locality information for each demo node. The input is a colon separated
+list of localities for each node. The i'th locality in the colon separated
+list sets the locality for the i'th demo cockroach node. For example:
 <PRE>
 
-	--demo-locality=region=us-east1,region=us-east2,region=us-east3
+--demo-locality=region=us-east1,az=1:region=us-east1,az=2:region=us-east1,az=3
 
-Assigns node 1's region to us-east1, node 2's region to us-east2,
-and node 3's region to us-east3.
+Assigns node1's region to us-east1 and availability zone to 1, node 2's
+region to us-east1 and availability zone to 2, and node 3's region
+to us-east1 and availability zone to 3.
 `,
 	}
 

pkg/cli/context.go

@@ -18,7 +18,6 @@ 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"
@@ -146,7 +145,7 @@ func initCLIDefaults() {
 
 	demoCtx.nodes = 1
 	demoCtx.useEmptyDatabase = false
-	demoCtx.localities = roachpb.Locality{}
+	demoCtx.localities = nil
 
 	initPreFlagsDefaults()
 
@@ -338,5 +337,5 @@ var sqlfmtCtx struct {
 var demoCtx struct {
 	nodes            int
 	useEmptyDatabase bool
-	localities       roachpb.Locality
+	localities       demoLocalityList
 }

pkg/cli/demo.go

@@ -18,7 +18,6 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/cli/cliflags"
-	"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"
@@ -91,14 +90,12 @@ func setupTransientServers(
 		return "", "", cleanup, errors.Errorf("must have a positive number of nodes")
 	}
 
-	var nodeLocalityTiers []roachpb.Tier
-	if len(demoCtx.localities.Tiers) != 0 {
+	if len(demoCtx.localities) != 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 {
+		if len(demoCtx.localities) != 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
@@ -135,8 +132,8 @@ func setupTransientServers(
 		if s != nil {
 			args.JoinAddr = s.ServingRPCAddr()
 		}
-		if nodeLocalityTiers != nil {
-			args.Locality = roachpb.Locality{Tiers: nodeLocalityTiers[i : i+1]}
+		if demoCtx.localities != nil {
+			args.Locality = demoCtx.localities[i]
 		}
 		serv := serverFactory.New(args).(*server.TestServer)
 		if err := serv.Start(args); err != nil {

pkg/cli/flags_util.go

@@ -37,7 +37,7 @@ var _ pflag.Value = &localityList{}
 // Type implements the pflag.Value interface.
 func (l *localityList) Type() string { return "localityList" }
 
-// String implements the pflag.Value interface.=
+// String implements the pflag.Value interface.
 func (l *localityList) String() string {
 	string := ""
 	for _, loc := range []roachpb.LocalityAddress(*l) {
@@ -47,7 +47,7 @@ func (l *localityList) String() string {
 	return string
 }
 
-// String implements the pflag.Value interface.
+// Set implements the pflag.Value interface.
 func (l *localityList) Set(value string) error {
 	*l = []roachpb.LocalityAddress{}
 
@@ -78,6 +78,35 @@ func (l *localityList) Set(value string) error {
 	return nil
 }
 
+// type used to implement parsing a list of localities for the cockroach demo command.
+type demoLocalityList []roachpb.Locality
+
+// Type implements the pflag.Value interface.
+func (l *demoLocalityList) Type() string { return "demoLocalityList" }
+
+// String implements the pflag.Value interface.
+func (l *demoLocalityList) String() string {
+	s := ""
+	for _, loc := range []roachpb.Locality(*l) {
+		s += loc.String()
+	}
+	return s
+}
+
+// Set implements the pflag.Value interface.
+func (l *demoLocalityList) Set(value string) error {
+	*l = []roachpb.Locality{}
+	locs := strings.Split(value, ":")
+	for _, value := range locs {
+		parsedLoc := &roachpb.Locality{}
+		if err := parsedLoc.Set(value); err != nil {
+			return err
+		}
+		*l = append(*l, *parsedLoc)
+	}
+	return nil
+}
+
 // This file contains definitions for data types suitable for use by
 // the flag+pflag packages.