Skip to content

Commit

Permalink
roachtest: clean up command-line flags
Browse files Browse the repository at this point in the history 
The code around command-line flags is pretty messy: they're defined in
many places; the name and description of a flag are far away from
the variable; the variable names look like local variables and in many
cases it's not obvious we're accessing a global.

This commit moves all flags to a separate subpackage,
`roachtestflags`, making all uses of global flags obvious. We also add
a bit of infrastructure to allow defining all information about a flag
right next to where the variable is declared.

We also provide a `Changed()` function that determines if a flag value
was changed (without the caller having to use the Command or even the
flag name).

There should be no functional changes (just some cosmetic improvements
to the flag usage texts).

Epic: none
Release note: None
  • Loading branch information
@RaduBerinde
RaduBerinde committed Oct 10, 2023
1 parent a34a352 commit e6b6b8c
Show file tree
Hide file tree
Showing 13 changed files with 795 additions and 354 deletions.
3 changes: 3 additions & 0 deletions pkg/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ ALL_TESTS = [
"//pkg/cmd/roachtest/clusterstats:clusterstats_test",
"//pkg/cmd/roachtest/option:option_test",
"//pkg/cmd/roachtest/registry:registry_test",
"//pkg/cmd/roachtest/roachtestflags:roachtestflags_test",
"//pkg/cmd/roachtest/roachtestutil/mixedversion:mixedversion_test",
"//pkg/cmd/roachtest/roachtestutil:roachtestutil_test",
"//pkg/cmd/roachtest/tests:tests_test",
Expand Down Expand Up @@ -1131,6 +1132,8 @@ GO_TARGETS = [
"//pkg/cmd/roachtest/option:option_test",
"//pkg/cmd/roachtest/registry:registry",
"//pkg/cmd/roachtest/registry:registry_test",
"//pkg/cmd/roachtest/roachtestflags:roachtestflags",
"//pkg/cmd/roachtest/roachtestflags:roachtestflags_test",
"//pkg/cmd/roachtest/roachtestutil/clusterupgrade:clusterupgrade",
"//pkg/cmd/roachtest/roachtestutil/mixedversion:mixedversion",
"//pkg/cmd/roachtest/roachtestutil/mixedversion:mixedversion_test",
Expand Down
3 changes: 2 additions & 1 deletion pkg/cmd/roachtest/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ go_library(
"//pkg/cmd/roachtest/cluster",
"//pkg/cmd/roachtest/option",
"//pkg/cmd/roachtest/registry",
"//pkg/cmd/roachtest/roachtestflags",
"//pkg/cmd/roachtest/roachtestutil",
"//pkg/cmd/roachtest/spec",
"//pkg/cmd/roachtest/test",
Expand Down Expand Up @@ -61,7 +62,6 @@ go_library(
"@com_github_prometheus_client_golang//prometheus/promhttp",
"@com_github_slack_go_slack//:slack",
"@com_github_spf13_cobra//:cobra",
"@com_github_spf13_pflag//:pflag",
"@org_golang_x_sync//errgroup",
],
)
Expand Down Expand Up @@ -93,6 +93,7 @@ go_test(
"//pkg/cmd/roachtest/cluster",
"//pkg/cmd/roachtest/option",
"//pkg/cmd/roachtest/registry",
"//pkg/cmd/roachtest/roachtestflags",
"//pkg/cmd/roachtest/spec",
"//pkg/cmd/roachtest/test",
"//pkg/internal/team",
Expand Down
130 changes: 40 additions & 90 deletions pkg/cmd/roachtest/cluster.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestflags"
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil"
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec"
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test"
Expand All @@ -51,71 +52,21 @@ import (
"github.com/cockroachdb/cockroach/pkg/util/timeutil"
"github.com/cockroachdb/errors"
_ "github.com/lib/pq"
"github.com/spf13/pflag"
)

func init() {
_ = roachprod.InitProviders()
}

var (
// user-specified path to crdb binary
cockroachPath string
// maps cpuArch to the corresponding crdb binary's absolute path
cockroach = make(map[vm.CPUArch]string)
// user-specified path to crdb binary with runtime assertions enabled (EA)
cockroachEAPath string
// maps cpuArch to the corresponding crdb binary with runtime assertions enabled (EA)
cockroachEA = make(map[vm.CPUArch]string)
// user-specified path to workload binary
workloadPath string
// maps cpuArch to the corresponding workload binary's absolute path
workload = make(map[vm.CPUArch]string)
// maps cpuArch to the corresponding dynamically-linked libraries' absolute paths
libraryFilePaths = make(map[vm.CPUArch][]string)
cloud = spec.GCE
// encryptionProbability controls when encryption-at-rest is enabled
// in a cluster for tests that have opted-in to metamorphic
// encryption (EncryptionMetamorphic).
//
// Tests that have opted-in to metamorphic encryption will run with
// encryption enabled by default (probability 1). In order to run
// them with encryption disabled (perhaps to reproduce a test
// failure), roachtest can be invoked with --metamorphic-encryption-probability=0
encryptionProbability float64
// Total probability with which new ARM64 clusters are provisioned, modulo test specs. which are incompatible.
// N.B. if all selected tests are incompatible with ARM64, then arm64Probability is effectively 0.
// In other words, ClusterSpec.Arch takes precedence over the arm64Probability flag.
arm64Probability float64
// Conditional probability with which new FIPS clusters are provisioned, modulo test specs. The total probability
// is the product of this and 1-arm64Probability.
// As in the case of arm64Probability, ClusterSpec.Arch takes precedence over the fipsProbability flag.
fipsProbability float64

instanceType string
localSSDArg bool
deprecatedRoachprodBinary string
// overrideOpts contains vm.CreateOpts override values passed from the cli.
overrideOpts vm.CreateOpts
// overrideFlagset represents the flags passed from the cli for
// `run` command (used to know if the value of a flag changed,
// for example: overrideFlagset("lifetime").Changed().
// TODO(ahmad/healthy-pod): extract runCmd (and other commands) from main
// to make it global and operate on runCmd.Flags() directly.
overrideFlagset *pflag.FlagSet
overrideNumNodes = -1
clusterName string
local bool
clusterWipe bool
zonesF string
teamCity bool
disableIssue bool
)

const (
defaultEncryptionProbability = 1
defaultFIPSProbability = 0
defaultARM64Probability = 0
)

type errBinaryOrLibraryNotFound struct {
Expand Down Expand Up @@ -178,7 +129,7 @@ func findBinary(

func findLibrary(libraryName string, os string, arch vm.CPUArch) (string, error) {
suffix := ".so"
if cloud == spec.Local {
if roachtestflags.Cloud == spec.Local {
switch runtime.GOOS {
case "linux":
case "freebsd":
Expand Down Expand Up @@ -308,9 +259,9 @@ func initBinariesAndLibraries() {
// see https://github.com/cockroachdb/cockroach/issues/104029.
defaultOSName := "linux"
defaultArch := vm.ArchAMD64
if cloud == spec.Local {
if roachtestflags.Cloud == spec.Local {
defaultOSName = runtime.GOOS
if arm64Probability == 1 {
if roachtestflags.ARM64Probability == 1 {
// N.B. if arm64Probability != 1, then we're running a local cluster with both arm64 and amd64.
defaultArch = vm.ArchARM64
}
Expand Down Expand Up @@ -361,14 +312,17 @@ func initBinariesAndLibraries() {
// We need to verify we have at least both the cockroach and the workload binaries.
var err error

cockroachPath := roachtestflags.CockroachPath
cockroachEAPath := roachtestflags.CockroachEAPath
workloadPath := roachtestflags.WorkloadPath
cockroach[defaultArch], _ = resolveBinary("cockroach", cockroachPath, defaultArch, true, false)
workload[defaultArch], _ = resolveBinary("workload", workloadPath, defaultArch, true, false)
cockroachEA[defaultArch], err = resolveBinary("cockroach-ea", cockroachEAPath, defaultArch, false, true)
if err != nil {
fmt.Fprintf(os.Stderr, "WARN: unable to find %q for %q: %s\n", "cockroach-ea", defaultArch, err)
}

if arm64Probability > 0 && defaultArch != vm.ArchARM64 {
if roachtestflags.ARM64Probability > 0 && defaultArch != vm.ArchARM64 {
fmt.Printf("Locating and verifying binaries for os=%q, arch=%q\n", defaultOSName, vm.ArchARM64)
// We need to verify we have all the required binaries for arm64.
cockroach[vm.ArchARM64], _ = resolveBinary("cockroach", cockroachPath, vm.ArchARM64, true, false)
Expand All @@ -378,7 +332,7 @@ func initBinariesAndLibraries() {
fmt.Fprintf(os.Stderr, "WARN: unable to find %q for %q: %s\n", "cockroach-ea", vm.ArchARM64, err)
}
}
if fipsProbability > 0 && defaultArch != vm.ArchFIPS {
if roachtestflags.FIPSProbability > 0 && defaultArch != vm.ArchFIPS {
fmt.Printf("Locating and verifying binaries for os=%q, arch=%q\n", defaultOSName, vm.ArchFIPS)
// We need to verify we have all the required binaries for fips.
cockroach[vm.ArchFIPS], _ = resolveBinary("cockroach", cockroachPath, vm.ArchFIPS, true, false)
Expand All @@ -392,11 +346,11 @@ func initBinariesAndLibraries() {
// In v20.2 or higher, optionally expect certain library files to exist.
// Since they may not be found in older versions, do not hard error if they are not found.
for _, arch := range []vm.CPUArch{vm.ArchAMD64, vm.ArchARM64, vm.ArchFIPS} {
if arm64Probability == 0 && defaultArch != vm.ArchARM64 && arch == vm.ArchARM64 {
if roachtestflags.ARM64Probability == 0 && defaultArch != vm.ArchARM64 && arch == vm.ArchARM64 {
// arm64 isn't used, skip finding libs for it.
continue
}
if fipsProbability == 0 && arch == vm.ArchFIPS {
if roachtestflags.FIPSProbability == 0 && arch == vm.ArchFIPS {
// fips isn't used, skip finding libs for it.
continue
}
Expand Down Expand Up @@ -838,29 +792,24 @@ func (f *clusterFactory) genName(cfg clusterConfig) string {
}

// createFlagsOverride updates opts with the override values passed from the cli.
func createFlagsOverride(flags *pflag.FlagSet, opts *vm.CreateOpts) {
if flags != nil {
if flags.Changed("lifetime") {
opts.Lifetime = overrideOpts.Lifetime
}
if flags.Changed("roachprod-local-ssd") {
opts.SSDOpts.UseLocalSSD = overrideOpts.SSDOpts.UseLocalSSD
}
if flags.Changed("filesystem") {
opts.SSDOpts.FileSystem = overrideOpts.SSDOpts.FileSystem
}
if flags.Changed("local-ssd-no-ext4-barrier") {
opts.SSDOpts.NoExt4Barrier = overrideOpts.SSDOpts.NoExt4Barrier
}
if flags.Changed("os-volume-size") {
opts.OsVolumeSize = overrideOpts.OsVolumeSize
}
if flags.Changed("geo") {
opts.GeoDistributed = overrideOpts.GeoDistributed
}
if flags.Changed("label") {
opts.CustomLabels = overrideOpts.CustomLabels
}
func createFlagsOverride(opts *vm.CreateOpts) {
if roachtestflags.Changed(&roachtestflags.Lifetime) {
opts.Lifetime = roachtestflags.Lifetime
}
if roachtestflags.Changed(&roachtestflags.OverrideUseLocalSSD) {
opts.SSDOpts.UseLocalSSD = roachtestflags.OverrideUseLocalSSD
}
if roachtestflags.Changed(&roachtestflags.OverrideFilesystem) {
opts.SSDOpts.FileSystem = roachtestflags.OverrideFilesystem
}
if roachtestflags.Changed(&roachtestflags.OverrideNoExt4Barrier) {
opts.SSDOpts.NoExt4Barrier = roachtestflags.OverrideNoExt4Barrier
}
if roachtestflags.Changed(&roachtestflags.OverrideOSVolumeSizeGB) {
opts.OsVolumeSize = roachtestflags.OverrideOSVolumeSizeGB
}
if roachtestflags.Changed(&roachtestflags.OverrideGeoDistributed) {
opts.GeoDistributed = roachtestflags.OverrideGeoDistributed
}
}

Expand All @@ -886,8 +835,8 @@ func (f *clusterFactory) newCluster(
return nil, nil, errors.Wrap(ctx.Err(), "newCluster")
}

if overrideFlagset != nil && overrideFlagset.Changed("nodes") {
cfg.spec.NodeCount = overrideNumNodes
if roachtestflags.Changed(&roachtestflags.OverrideNumNodes) {
cfg.spec.NodeCount = roachtestflags.OverrideNumNodes
}

if cfg.spec.NodeCount == 0 {
Expand Down Expand Up @@ -921,9 +870,9 @@ func (f *clusterFactory) newCluster(
UseIOBarrierOnLocalSSD: cfg.useIOBarrier,
PreferredArch: cfg.arch,
}
params.Defaults.MachineType = instanceType
params.Defaults.Zones = zonesF
params.Defaults.PreferLocalSSD = localSSDArg
params.Defaults.MachineType = roachtestflags.InstanceType
params.Defaults.Zones = roachtestflags.Zones
params.Defaults.PreferLocalSSD = roachtestflags.PreferLocalSSD

// The ClusterName is set below in the retry loop to ensure
// that each create attempt gets a unique cluster name.
Expand All @@ -938,7 +887,8 @@ func (f *clusterFactory) newCluster(
providerOptsContainer.SetProviderOpts(cfg.spec.Cloud, providerOpts)
}

createFlagsOverride(overrideFlagset, &createVMOpts)
createFlagsOverride(&createVMOpts)

// Make sure expiration is changed if --lifetime override flag
// is passed.
cfg.spec.Lifetime = createVMOpts.Lifetime
Expand Down Expand Up @@ -1070,7 +1020,7 @@ func attachToExistingCluster(
return nil, err
}
if !opt.skipWipe {
if clusterWipe {
if roachtestflags.ClusterWipe {
if err := c.WipeE(ctx, l, false /* preserveCerts */, c.All()); err != nil {
return nil, err
}
Expand Down Expand Up @@ -1717,7 +1667,7 @@ func (c *clusterImpl) doDestroy(ctx context.Context, l *logger.Logger) <-chan st
return inFlight
}

if clusterWipe {
if roachtestflags.ClusterWipe {
if c.destroyState.owned {
l.PrintfCtx(ctx, "destroying cluster %s...", c)
c.status("destroying cluster")
Expand Down Expand Up @@ -2695,9 +2645,9 @@ func archForTest(ctx context.Context, l *logger.Logger, testSpec registry.TestSp
// CPU architecture is unspecified, choose one according to the
// probability distribution.
var arch vm.CPUArch
if prng.Float64() < arm64Probability {
if prng.Float64() < roachtestflags.ARM64Probability {
arch = vm.ArchARM64
} else if prng.Float64() < fipsProbability {
} else if prng.Float64() < roachtestflags.FIPSProbability {
// N.B. branch is taken with probability
// (1 - arm64Probability) * fipsProbability
// which is P(fips & amd64).
Expand Down
Loading

0 comments on commit e6b6b8c

Please sign in to comment.