Skip to content

Commit

Permalink
Merge pull request #114432 from RaduBerinde/backport23.2-112078
Browse files Browse the repository at this point in the history 
release-23.2: roachtest: clean up command-line flags
  • Loading branch information
@RaduBerinde
RaduBerinde authored Nov 15, 2023
2 parents 0cd5220 + a357561 commit 6f8997e
Show file tree
Hide file tree
Showing 13 changed files with 808 additions and 362 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 @@ -1129,6 +1130,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 @@ -94,6 +94,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: 39 additions & 91 deletions pkg/cmd/roachtest/cluster.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,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 @@ -53,71 +54,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 @@ -180,7 +131,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 @@ -310,9 +261,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 @@ -363,14 +314,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 @@ -380,7 +334,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 @@ -394,11 +348,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 @@ -845,29 +799,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 @@ -893,8 +842,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 @@ -928,9 +877,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 @@ -948,8 +897,7 @@ func (f *clusterFactory) newCluster(
if cfg.spec.UbuntuVersion.IsOverridden() {
createVMOpts.UbuntuVersion = cfg.spec.UbuntuVersion
}

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 @@ -1081,7 +1029,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 @@ -1728,7 +1676,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 @@ -2760,9 +2708,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 6f8997e

Please sign in to comment.