Skip to content

Commit

Permalink
demo: Add a flag to automatically apply geo-partitioned replicas.
Browse files Browse the repository at this point in the history
Work for cockroachdb#39945.

This PR adds a Partition function to the workload.Hooks struct so that
workloads that have a partitioning step can just implement that function
within the workload package.

Release note (cli change): Add an option for cockroach demo to
automatically apply the geo-partitioned replicas topology to the movr
dataset using the --geo-partitioned-replicas flag.
  • Loading branch information
rohany authored and ajwerner committed Sep 13, 2019
1 parent e184e75 commit 2ec3b13
Show file tree
Hide file tree
Showing 9 changed files with 301 additions and 16 deletions.
1 change: 1 addition & 0 deletions pkg/ccl/cliccl/demo.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ func getLicense(clusterID uuid.UUID) (string, error) {
return "", err
}
defer resp.Body.Close()

if resp.StatusCode != http.StatusOK {
return "", errors.New("unable to connect to licensing endpoint")
}
Expand Down
2 changes: 1 addition & 1 deletion pkg/ccl/workloadccl/allccl/all_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -251,7 +251,7 @@ func TestDeterministicInitialData(t *testing.T) {
`intro`: 0x81c6a8cfd9c3452a,
`json`: 0xcbf29ce484222325,
`ledger`: 0xebe27d872d980271,
`movr`: 0x4f19a54c7e779f9c,
`movr`: 0x6a094e9d15a07970,
`queue`: 0xcbf29ce484222325,
`rand`: 0xcbf29ce484222325,
`roachmart`: 0xda5e73423dbdb2d9,
Expand Down
12 changes: 12 additions & 0 deletions pkg/cli/cliflags/flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -918,6 +918,18 @@ to us-east1 and availability zone to 3.
`,
}

DemoGeoPartitionedReplicas = FlagInfo{
Name: "geo-partitioned-replicas",
Description: `
When used with the Movr dataset, create a 9 node cluster and automatically apply
the geo-partitioned replicas topology across 3 virtual regions named us-east1, us-west1, and
europe-west1. This command will fail with an error if an enterprise license could not
be acquired, or if the Movr dataset is not used. More information about the geo-partitioned
replicas topology can be found at this URL:
https://www.cockroachlabs.com/docs/v19.1/topology-geo-partitioned-replicas.html
`,
}

UseEmptyDatabase = FlagInfo{
Name: "empty",
Description: `
Expand Down
10 changes: 6 additions & 4 deletions pkg/cli/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ func initCLIDefaults() {
demoCtx.useEmptyDatabase = false
demoCtx.runWorkload = false
demoCtx.localities = nil
demoCtx.geoPartitionedReplicas = false

initPreFlagsDefaults()

Expand Down Expand Up @@ -336,8 +337,9 @@ var sqlfmtCtx struct {
// demoCtx captures the command-line parameters of the `demo` command.
// Defaults set by InitCLIDefaults() above.
var demoCtx struct {
nodes int
useEmptyDatabase bool
runWorkload bool
localities demoLocalityList
nodes int
useEmptyDatabase bool
runWorkload bool
localities demoLocalityList
geoPartitionedReplicas bool
}
96 changes: 89 additions & 7 deletions pkg/cli/demo.go
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,10 @@ subcommands: e.g. "cockroach demo startrek". See --help for a full list.
By default, the 'movr' dataset is pre-loaded. You can also use --empty
to avoid pre-loading a dataset.
cockroach demo attempts to connect to a Cockroach Labs server to obtain a
temporary enterprise license for demoing enterprise features and enable
telemetry back to Cockroach Labs. In order to disable this behavior, set the
environment variable "COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING".
cockroach demo attempts to connect to a Cockroach Labs server to obtain a
temporary enterprise license for demoing enterprise features and enable
telemetry back to Cockroach Labs. In order to disable this behavior, set the
environment variable "COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING" to true.
`,
Example: ` cockroach demo`,
Args: cobra.NoArgs,
Expand Down Expand Up @@ -207,6 +207,10 @@ func setupTransientServers(
}
urlStr := url.String()

// Communicate information about license acquisition to services
// that depend on it.
licenseSuccess := make(chan bool, 1)

// Start up the update check loop.
// We don't do this in (*server.Server).Start() because we don't want it
// in tests.
Expand All @@ -230,8 +234,13 @@ func setupTransientServers(
const msg = "Unable to acquire demo license. Enterprise features are not enabled in this session.\n"
fmt.Fprint(stderr, msg)
}
licenseSuccess <- success
}()
}
} else {
// If we aren't supposed to check for a license, then automatically
// notify failure.
licenseSuccess <- false
}

// If there is a load generator, create its database and load its
Expand All @@ -253,10 +262,40 @@ func setupTransientServers(
return ``, ``, cleanup, err
}

partitioningComplete := make(chan struct{}, 1)
// If we are requested to prepartition our data spawn a goroutine to do the partitioning.
if demoCtx.geoPartitionedReplicas {
go func() {
success := <-licenseSuccess
// Only try partitioning if license acquisition was successful.
if success {
db, err := gosql.Open("postgres", urlStr)
if err != nil {
exitWithError("demo", err)
}
defer db.Close()
// Based on validation done in setup, we know that this workload has a partitioning step.
if err := gen.(workload.Hookser).Hooks().Partition(db); err != nil {
exitWithError("demo", err)
}
partitioningComplete <- struct{}{}
} else {
const msg = "license acquisition was unsuccessful. Enterprise features are needed to partition data"
exitWithError("demo", errors.New(msg))
}
}()
}

if demoCtx.runWorkload {
if err := runWorkload(ctx, gen, urlStr, stopper); err != nil {
return ``, ``, cleanup, err
}
go func() {
// If partitioning was requested, wait for that to complete before running the workload.
if demoCtx.geoPartitionedReplicas {
<-partitioningComplete
}
if err := runWorkload(ctx, gen, urlStr, stopper); err != nil {
exitWithError("demo", err)
}
}()
}
}

Expand Down Expand Up @@ -319,6 +358,49 @@ func runDemo(cmd *cobra.Command, gen workload.Generator) error {
return errors.New("cannot run a workload against an empty database")
}

// Make sure that the user didn't request to have a topology and an empty database.
if demoCtx.geoPartitionedReplicas && demoCtx.useEmptyDatabase {
return errors.New("cannot setup geo-partitioned replicas topology on an empty database")
}

// Make sure that the Movr database is selected when automatically partitioning.
if demoCtx.geoPartitionedReplicas && (gen == nil || gen.Meta().Name != "movr") {
return errors.New("--geo-partitioned-replicas must be used with the Movr dataset")
}

// If the geo-partitioned replicas flag was given and the demo localities have changed, throw an error.
if demoCtx.geoPartitionedReplicas && demoCtx.localities != nil {
return errors.New("--demo-locality cannot be used with --geo-partitioned-replicas")
}

// If the geo-partitioned replicas flag was given and the nodes have changed, throw an error.
if demoCtx.geoPartitionedReplicas && cmd.Flags().Lookup(cliflags.DemoNodes.Name).Changed {
return errors.New("--nodes cannot be used with --geo-partitioned-replicas")
}

// If geo-partition-replicas is requested, make sure the workload has a Partitioning step.
if demoCtx.geoPartitionedReplicas {
configErr := errors.New(fmt.Sprintf("workload %s is not configured to have a partitioning step", gen.Meta().Name))
hookser, ok := gen.(workload.Hookser)
if !ok {
return configErr
}
if hookser.Hooks().Partition == nil {
return configErr
}
}

// Th geo-partitioned replicas demo only works on a 9 node cluster, so set the node count as such.
// Ignore input user localities so that the nodes have the same attributes/localities as expected.
if demoCtx.geoPartitionedReplicas {
const msg = `#
# --geo-partitioned replicas operates on a 9 node cluster.
# The cluster size has been changed from the default to 9 nodes.`
fmt.Println(msg)
demoCtx.nodes = 9
demoCtx.localities = nil
}

connURL, adminURL, cleanup, err := setupTransientServers(cmd, gen)
defer cleanup()
if err != nil {
Expand Down
1 change: 1 addition & 0 deletions pkg/cli/flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -586,6 +586,7 @@ func init() {
IntFlag(demoFlags, &demoCtx.nodes, cliflags.DemoNodes, 1)
BoolFlag(demoFlags, &demoCtx.runWorkload, cliflags.RunDemoWorkload, false)
VarFlag(demoFlags, &demoCtx.localities, cliflags.DemoNodeLocality)
BoolFlag(demoFlags, &demoCtx.geoPartitionedReplicas, cliflags.DemoGeoPartitionedReplicas, false)
// 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)
Expand Down
87 changes: 87 additions & 0 deletions pkg/cli/interactive_tests/test_demo_partitioning.tcl
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
#! /usr/bin/env expect -f

source [file join [file dirname $argv0] common.tcl]

start_test "Expect partitioning succeeds"
# test that partitioning works if a license could be acquired
spawn $argv demo --geo-partitioned-replicas

# wait for the shell to start up
eexpect "movr>"

# send multiple "SHOW PARTITIONS" requests to the DB as partitioning is happen asynchronously.
for {set i 0} {$i < 10} {incr i} {
send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\];\r"
sleep 1
}

# The number of partitions across the MovR database we expect is 24.
eexpect "24"
eexpect "(1 row)"
eexpect "movr>"

send "SHOW PARTITIONS FROM TABLE vehicles;\r"

# Verify the partitions are as we expect
send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\] WHERE partition_name='us_west';\r"
eexpect "8"
eexpect "(1 row)"
eexpect "movr>"

send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\] WHERE partition_name='us_west' AND partition_value='(''seattle''), (''san francisco''), (''los angeles'')';\r"
eexpect "8"
eexpect "(1 row)"
eexpect "movr>"

send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\] WHERE partition_name='us_west' AND zone_config='constraints = ''\[+region=us-west1\]''';\r"
eexpect "8"
eexpect "(1 row)"
eexpect "movr>"

send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\] WHERE partition_name='us_east';\r"
eexpect "8"
eexpect "(1 row)"
eexpect "movr>"

send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\] WHERE partition_name='us_east' AND partition_value='(''new york''), (''boston''), (''washington dc'')';\r"
eexpect "8"
eexpect "(1 row)"
eexpect "movr>"

send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\] WHERE partition_name='us_east' AND zone_config='constraints = ''\[+region=us-east1\]''';\r"
eexpect "8"
eexpect "(1 row)"
eexpect "movr>"

send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\] WHERE partition_name='europe_west';\r"
eexpect "8"
eexpect "(1 row)"
eexpect "movr>"

send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\] WHERE partition_name='europe_west' AND partition_value='(''amsterdam''), (''paris''), (''rome'')';\r"
eexpect "8"
eexpect "(1 row)"
eexpect "movr>"

send "SELECT count(*) FROM \[SHOW PARTITIONS FROM DATABASE movr\] WHERE partition_name='europe_west' AND zone_config='constraints = ''\[+region=europe-west1\]''';\r"
eexpect "8"
eexpect "(1 row)"
eexpect "movr>"

interrupt
eexpect eof
end_test


start_test "Expect an error if geo-partitioning is requested and a license cannot be acquired"

# set the proper environment variable
set env(COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING) "true"
spawn $argv demo --geo-partitioned-replicas
# expect a failure
eexpect "Error: license acquisition was unsuccessful. Enterprise features are needed to partition data"
# clean up after the test
interrupt
eexpect eof
end_test

Loading

0 comments on commit 2ec3b13

Please sign in to comment.