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

[rohany]

The demo command now accepts a --demo-locality= flag,
which sets the locality of the i'th node in the demo to the
i'th locality setting input by the user.

Fixes #39450.

Release note (cli change): Add an option to set node
localities in cockroach demo.

[cockroach-teamcity]

This change is 

[rohany]

./cockroach demo --nodes 2 --demo-locality=region=us-east1,region=us-east2

[email protected]:57991/movr> create table t (x int primary key);
CREATE TABLE

[email protected]:57991/movr>
[email protected]:57991/movr> alter table t partition by list (x) ( partition p1 values in (1), partition p2 values in (2));
ALTER TABLE

[email protected]:57991/movr> alter partition p1 of table t configure zone using constraints='[+region=us-east1]';
CONFIGURE ZONE 1

[email protected]:57991/movr> alter partition p2 of table t configure zone using constraints='[+region=us-east2]';
CONFIGURE ZONE 1

[email protected]:57991/movr> show ranges from table t;
  start_key | end_key | range_id | replicas | lease_holder |    locality
+-----------+---------+----------+----------+--------------+-----------------+
  NULL      | /1      |       27 | {1}      |            1 | region=us-east1
  /1        | /2      |       28 | {1}      |            1 | region=us-east1
  /2        | /3      |       29 | {2}      |            2 | region=us-east2
  /3        | NULL    |       30 | {1,2}    |            2 | region=us-east2
(4 rows)

[jordanlewis]

nit: the trailing period thing :)

[knz]

The basic idea and implementation is good but since you're in this area you may as well refactor/simplify this code and add a missing sanity check. See my suggestions below.

Reviewed 4 of 4 files at r1.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @ajwerner, @awoods187, @jordanlewis, and @rohany)

---

pkg/cli/demo.go, line 136 at r1 (raw file):

	if len(demoCtx.localities.Tiers) != 0 {
		if len(demoCtx.localities.Tiers) != demoCtx.nodes {
			return "", "", cleanup, errors.Errorf("number of localities specified must equal number of nodes.")

nit: no trailing period in message

---

pkg/cli/demo.go, line 138 at r1 (raw file):

			return "", "", cleanup, errors.Errorf("number of localities specified must equal number of nodes.")
		}
		args.Locality = roachpb.Locality{Tiers: []roachpb.Tier{demoCtx.localities.Tiers[0]}}

Remove this line and ...

---

pkg/cli/demo.go, line 141 at r1 (raw file):

	}

	serverFactory := server.TestServerFactory

Move this line and the 4 that follow in the loop below.

---

pkg/cli/demo.go, line 146 at r1 (raw file):

		return connURL, adminURL, cleanup, err
	}
	args.JoinAddr = s.ServingAddr()

As you move this line into the loop, surround it with if i > 0 { ... }
Also for i == 0 capture the created server so it's available under the loop.

Finally, add a check that demoCtx.nodes > 0 above.

---

pkg/cli/demo.go, line 147 at r1 (raw file):

	}
	args.JoinAddr = s.ServingAddr()
	for i := 0; i < demoCtx.nodes-1; i++ {

Change to iterate until demoCtx.nodes without -1

[rohany]

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @ajwerner, @awoods187, @jordanlewis, and @knz)

---

pkg/cli/demo.go, line 138 at r1 (raw file):

Previously, knz (kena) wrote…

Remove this line and ...

Done.

---

pkg/cli/demo.go, line 141 at r1 (raw file):

Previously, knz (kena) wrote…

Move this line and the 4 that follow in the loop below.

Done.

---

pkg/cli/demo.go, line 146 at r1 (raw file):

Previously, knz (kena) wrote…

As you move this line into the loop, surround it with if i > 0 { ... }
Also for i == 0 capture the created server so it's available under the loop.

Finally, add a check that demoCtx.nodes > 0 above.

Done.

---

pkg/cli/demo.go, line 147 at r1 (raw file):

Previously, knz (kena) wrote…

Change to iterate until demoCtx.nodes without -1

Done.

[rohany]

I restructured the code. RFAL

[knz]

thanks just a few minor tweaks and this is good to go

Reviewed 2 of 4 files at r2.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @ajwerner, @awoods187, @jordanlewis, @knz, and @rohany)

---

pkg/cli/context.go, line 338 at r2 (raw file):

	nodes            int
	useEmptyDatabase bool
	localities       roachpb.Locality

I forgot to ask, did you think of resetting the field in InitCLIDefaults

---

pkg/cli/demo.go, line 151 at r2 (raw file):

	for i := 0; i < demoCtx.nodes; i++ {
		if len(demoCtx.localities.Tiers) != 0 {
			args.Locality = roachpb.Locality{Tiers: []roachpb.Tier{demoCtx.localities.Tiers[i]}}

reset args.Locality to nil otherwise

---

pkg/cli/demo.go, line 159 at r2 (raw file):

		if i == 0 {
			s = serv
			args.JoinAddr = s.ServingRPCAddr()

} else { args.JoinAddr = s. ...

(join is set for every node i > 1)

[rohany]

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @ajwerner, @awoods187, @jordanlewis, and @knz)

---

pkg/cli/context.go, line 338 at r2 (raw file):

Previously, knz (kena) wrote…

I forgot to ask, did you think of resetting the field in InitCLIDefaults

Done.

---

pkg/cli/demo.go, line 151 at r2 (raw file):

Previously, knz (kena) wrote…

reset args.Locality to nil otherwise

The default value is an empty roachpb.Locality

---

pkg/cli/demo.go, line 159 at r2 (raw file):

Previously, knz (kena) wrote…

} else { args.JoinAddr = s. ...

(join is set for every node i > 1)

Before this patch args.JoinAddr was set only once to the address of the first server created (all other nodes share the same args object) -- this does the same thing.

[rohany]

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @ajwerner, @awoods187, @jordanlewis, and @knz)

---

pkg/cli/demo.go, line 151 at r2 (raw file):

Previously, rohany (Rohan Yadav) wrote…

The default value is an empty roachpb.Locality

Sorry, I misunderstood -- args.Locality is already the default locality value, and is only changed when it is possible to assign all nodes a locality -- so we never get into a "bad" state here.

[knz]

Reviewed 2 of 4 files at r2, 1 of 1 files at r3.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @ajwerner, @awoods187, @jordanlewis, and @rohany)

---

pkg/cli/demo.go, line 151 at r2 (raw file):

Previously, rohany (Rohan Yadav) wrote…

Sorry, I misunderstood -- args.Locality is already the default locality value, and is only changed when it is possible to assign all nodes a locality -- so we never get into a "bad" state here.

Oh now I understand. Thanks for explaining.

But the check on the length demoCtx.localities.Tiers is not trivial to understand. At least it deserves a comment. Perhaps it would be better to do something like this:

• above:

var localityTiers []roachpb.Tier
if len(...) > 0 {
   ...
   localityTiers = demoCtx.localities.Tiers

• here:

if localityTiers != 0 {
   args.Locality.Tiers = localityTiers[i:i+1]
}

(Also note how you can reuse the roachpb.Tier slice by re-slicing it every time. Better than creating a new slice for each server.)

---

pkg/cli/demo.go, line 159 at r2 (raw file):

Previously, rohany (Rohan Yadav) wrote…

Before this patch args.JoinAddr was set only once to the address of the first server created (all other nodes share the same args object) -- this does the same thing.

yes I understand this does the same thing. Do you find it equally readable thought? I always find it confusing when the end of the loop body sets some variables for the next iteration, especially in this case "the first iteration sets the variable for every next iteration".

In general I like my loop bodies to be dependent only on the loop counter. This makes it possible to consider each iteration independently from the others.

I'll leave it to you though. that, or at the very least an explanatory comment.

[knz]

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @ajwerner, @awoods187, @jordanlewis, and @rohany)

---

pkg/cli/demo.go, line 151 at r2 (raw file):

Previously, knz (kena) wrote…

Oh now I understand. Thanks for explaining.

But the check on the length demoCtx.localities.Tiers is not trivial to understand. At least it deserves a comment. Perhaps it would be better to do something like this:

• above:

var localityTiers []roachpb.Tier
if len(...) > 0 {
   ...
   localityTiers = demoCtx.localities.Tiers

• here:

if localityTiers != 0 {
   args.Locality.Tiers = localityTiers[i:i+1]
}

(Also note how you can reuse the roachpb.Tier slice by re-slicing it every time. Better than creating a new slice for each server.)

I meant != nil but you get the drift

[rohany]

Alright, I cleaned up the code based on your suggestions -- RFAL

[knz]

with nit

Reviewed 1 of 1 files at r4.
Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @ajwerner, @awoods187, @jordanlewis, and @rohany)

---

pkg/cli/demo.go, line 159 at r2 (raw file):

Previously, knz (kena) wrote…

yes I understand this does the same thing. Do you find it equally readable thought? I always find it confusing when the end of the loop body sets some variables for the next iteration, especially in this case "the first iteration sets the variable for every next iteration".

In general I like my loop bodies to be dependent only on the loop counter. This makes it possible to consider each iteration independently from the others.

I'll leave it to you though. that, or at the very least an explanatory comment.

If you choose (as you did here) to keep the assignment done on every iteration, for the benefit of the next iteration (as opposed to an assignment at the beginning of the loop body for the benefit of the current iteration), I still think this deserves an explanatory comment. Your future self will thank you for it.

[rohany]

Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @ajwerner, @awoods187, @jordanlewis, and @knz)

---

pkg/cli/demo.go, line 151 at r2 (raw file):

Previously, knz (kena) wrote…

I meant != nil but you get the drift

Done.

---

pkg/cli/demo.go, line 159 at r2 (raw file):

Previously, knz (kena) wrote…

If you choose (as you did here) to keep the assignment done on every iteration, for the benefit of the next iteration (as opposed to an assignment at the beginning of the loop body for the benefit of the current iteration), I still think this deserves an explanatory comment. Your future self will thank you for it.

sorry it took me so long -- I understand now why you don't like the current version of the code. I agree that it is not easy to reason about. I changed it to be clearer -- setting the args at the beginning of the loop iteration.

[rohany]

bors r+

[craig]

Build succeeded

• GitHub CI (Cockroach)

[andreimatei]

This is a fantastic feature that I've just tried to use, but I believe what's going on here doesn't quite compute. We're taking a comma separated value from the command line, we parse it into a roachpb.Locality which conveniently implements the flag interface - and so the different values between the commas represent different locality "tiers" - and then we take that one locality and we extract each tier to serve as the single tier of a new locality for each node...

This is a slightly outrageous abuse to parse a list of localities as a single locality, and also this doesn't let me specify localities with multiple tiers for my demo nodes.

Can I kindly ask for an improvement? :P

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @ajwerner, @awoods187, @jordanlewis, and @knz)

[rohany]

It is an abuse to reuse the Locality tiers as a list of single tiered localities instead :) -- How do you propose an interface for setting multiple tiered localities for the demo nodes would work?

Something like --

./cockroach demo --nodes=3 --demo-localities=region=us-east1,az=1;region=us-east-1,az=2;region=us-east-1,az=3

i.e using semi colons to separate the tiers of localities.

[knz]

Semicolons are not shell friendly. Suggestion to use colons instead. -- Sent from my Android device with K-9 Mail. Please excuse my brevity.

[rohany]

Colons make sense -- I can start an issue tracking this.