roachtest: clean up zones specification

`ClusterSpec` now provides options for GCE and AWS zones
specification.

Epic: none
Release note: None

pkg/cmd/roachtest/spec/cluster_spec.go

@@ -71,6 +71,10 @@ type ClusterSpec struct {
 	// the spec.
 	defaultInstanceType string
 
+	// TODO(radu): defaultZones is the default zones specification (unless
+	// overridden by GCE.Zones or AWS.Zones); it does not belong in the spec.
+	defaultZones string
+
 	Arch      vm.CPUArch // CPU architecture; auto-chosen if left empty
 	NodeCount int
 	// CPUs is the number of CPUs per node.
@@ -80,7 +84,6 @@ type ClusterSpec struct {
 	RAID0                bool
 	VolumeSize           int
 	PreferLocalSSD       bool
-	Zones                string
 	Geo                  bool
 	Lifetime             time.Duration
 	ReusePolicy          clusterReusePolicy
@@ -99,13 +102,15 @@ type ClusterSpec struct {
 		MachineType    string
 		MinCPUPlatform string
 		VolumeType     string
+		Zones          string
 	}
 
 	// AWS-specific arguments. These values apply only on clusters instantiated on AWS.
 	AWS struct {
 		MachineType string
 		// VolumeThroughput is the min provisioned EBS volume throughput.
 		VolumeThroughput int
+		Zones            string
 	}
 }
 
@@ -331,9 +336,21 @@ func (s *ClusterSpec) RoachprodOpts(
 			createVMOpts.SSDOpts.FileSystem = vm.Zfs
 		}
 	}
+
+	zonesStr := s.defaultZones
+	switch s.Cloud {
+	case AWS:
+		if s.AWS.Zones != "" {
+			zonesStr = s.AWS.Zones
+		}
+	case GCE:
+		if s.GCE.Zones != "" {
+			zonesStr = s.GCE.Zones
+		}
+	}
 	var zones []string
-	if s.Zones != "" {
-		zones = strings.Split(s.Zones, ",")
+	if zonesStr != "" {
+		zones = strings.Split(zonesStr, ",")
 		if !s.Geo {
 			zones = zones[:1]
 		}

pkg/cmd/roachtest/spec/option.go

@@ -80,12 +80,10 @@ func Geo() Option {
 	}
 }
 
-// Zones is a node option which requests Geo-distributed nodes. Note that this
-// overrides the --zones flag and is useful for tests that require running on
-// specific Zones.
-func Zones(zones string) Option {
+// DefaultZones sets the default zones (set with the --zones flag).
+func DefaultZones(zones string) Option {
 	return func(spec *ClusterSpec) {
-		spec.Zones = zones
+		spec.defaultZones = zones
 	}
 }
 
@@ -219,6 +217,17 @@ func GCEVolumeType(volumeType string) Option {
 	}
 }
 
+// GCEZones is a node option which requests Geo-distributed nodes; only applies
+// when the test runs on GCE.
+//
+// Note that this overrides the --zones flag and is useful for tests that
+// require running on specific zones.
+func GCEZones(zones string) Option {
+	return func(spec *ClusterSpec) {
+		spec.GCE.Zones = zones
+	}
+}
+
 // AWSMachineType sets the machine (instance) type when the cluster is on AWS.
 func AWSMachineType(machineType string) Option {
 	return func(spec *ClusterSpec) {
@@ -233,3 +242,14 @@ func AWSVolumeThroughput(throughput int) Option {
 		spec.AWS.VolumeThroughput = throughput
 	}
 }
+
+// AWSZones is a node option which requests Geo-distributed nodes; only applies
+// when the test runs on AWS.
+//
+// Note that this overrides the --zones flag and is useful for tests that
+// require running on specific zones.
+func AWSZones(zones string) Option {
+	return func(spec *ClusterSpec) {
+		spec.AWS.Zones = zones
+	}
+}

pkg/cmd/roachtest/test_registry.go

@@ -100,7 +100,7 @@ func (r *testRegistryImpl) MakeClusterSpec(nodeCount int, opts ...spec.Option) s
 		finalOpts = append(finalOpts, spec.PreferLocalSSD(true))
 	}
 	if r.zones != "" {
-		finalOpts = append(finalOpts, spec.Zones(r.zones))
+		finalOpts = append(finalOpts, spec.DefaultZones(r.zones))
 	}
 	finalOpts = append(finalOpts, opts...)
 	return spec.MakeClusterSpec(r.cloud, r.instanceType, nodeCount, finalOpts...)

pkg/cmd/roachtest/test_registry_test.go

@@ -29,11 +29,10 @@ func TestMakeTestRegistry(t *testing.T) {
 		require.Equal(t, "foo", r.instanceType)
 		require.Equal(t, spec.AWS, r.cloud)
 
-		s := r.MakeClusterSpec(100, spec.Geo(), spec.Zones("zone99"), spec.CPU(12),
+		s := r.MakeClusterSpec(100, spec.Geo(), spec.CPU(12),
 			spec.PreferLocalSSD(true))
 		require.EqualValues(t, 100, s.NodeCount)
 		require.True(t, s.Geo)
-		require.Equal(t, "zone99", s.Zones)
 		require.EqualValues(t, 12, s.CPUs)
 		require.True(t, s.PreferLocalSSD)
 

pkg/cmd/roachtest/tests/admission_control_database_drop.go

@@ -31,12 +31,12 @@ func registerDatabaseDrop(r registry.Registry) {
 	clusterSpec := r.MakeClusterSpec(
 		10, /* nodeCount */
 		spec.CPU(8),
-		spec.Zones("us-east1-b"),
 		spec.VolumeSize(500),
 		spec.Cloud(spec.GCE),
 		spec.GCEMinCPUPlatform("Intel Ice Lake"),
 		spec.GCEVolumeType("pd-ssd"),
 		spec.GCEMachineType("n2-standard-8"),
+		spec.GCEZones("us-east1-b"),
 	)
 
 	r.Add(registry.TestSpec{

pkg/cmd/roachtest/tests/admission_control_index_backfill.go

@@ -32,12 +32,12 @@ func registerIndexBackfill(r registry.Registry) {
 	clusterSpec := r.MakeClusterSpec(
 		10, /* nodeCount */
 		spec.CPU(8),
-		spec.Zones("us-east1-b"),
 		spec.VolumeSize(500),
 		spec.Cloud(spec.GCE),
 		spec.GCEMinCPUPlatform("Intel Ice Lake"),
 		spec.GCEVolumeType("pd-ssd"),
 		spec.GCEMachineType("n2-standard-8"),
+		spec.GCEZones("us-east1-b"),
 	)
 
 	r.Add(registry.TestSpec{

pkg/cmd/roachtest/tests/cluster_to_cluster.go

@@ -1002,7 +1002,7 @@ func c2cRegisterWrapper(
 		allZones = append(allZones, sp.multiregion.srcLocalities...)
 		allZones = append(allZones, sp.multiregion.destLocalities...)
 		allZones = append(allZones, sp.multiregion.workloadNodeZone)
-		clusterOps = append(clusterOps, spec.Zones(strings.Join(allZones, ",")))
+		clusterOps = append(clusterOps, spec.GCEZones(strings.Join(allZones, ",")))
 		clusterOps = append(clusterOps, spec.Geo())
 	}
 
@@ -1175,7 +1175,7 @@ func registerClusterToCluster(r registry.Registry) {
 				destLocalities:   []string{"us-central1-b", "us-west1-b", "us-west1-b", "us-west1-b"},
 				workloadNodeZone: "us-west1-b",
 			},
-			clouds: registry.AllExceptAWS,
+			clouds: registry.OnlyGCE,
 			suites: registry.Suites("nightly"),
 		},
 		{

pkg/cmd/roachtest/tests/connection_latency.go

@@ -122,8 +122,8 @@ func registerConnectionLatencyTest(r registry.Registry) {
 		Owner:     registry.OwnerSQLFoundations,
 		Benchmark: true,
 		// Add one more node for load node.
-		Cluster:          r.MakeClusterSpec(numNodes+1, spec.Zones(regionUsCentral)),
-		CompatibleClouds: registry.AllExceptAWS,
+		Cluster:          r.MakeClusterSpec(numNodes+1, spec.GCEZones(regionUsCentral)),
+		CompatibleClouds: registry.OnlyGCE,
 		Suites:           registry.Suites(registry.Nightly),
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
 			runConnectionLatencyTest(ctx, t, c, numNodes, 1, false /*password*/)
@@ -140,8 +140,8 @@ func registerConnectionLatencyTest(r registry.Registry) {
 		Name:             fmt.Sprintf("connection_latency/nodes=%d/multiregion/certs", numMultiRegionNodes),
 		Owner:            registry.OwnerSQLFoundations,
 		Benchmark:        true,
-		Cluster:          r.MakeClusterSpec(numMultiRegionNodes+loadNodes, spec.Geo(), spec.Zones(geoZonesStr)),
-		CompatibleClouds: registry.AllExceptAWS,
+		Cluster:          r.MakeClusterSpec(numMultiRegionNodes+loadNodes, spec.Geo(), spec.GCEZones(geoZonesStr)),
+		CompatibleClouds: registry.OnlyGCE,
 		Suites:           registry.Suites(registry.Nightly),
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
 			runConnectionLatencyTest(ctx, t, c, numMultiRegionNodes, numZones, false /*password*/)
@@ -152,8 +152,8 @@ func registerConnectionLatencyTest(r registry.Registry) {
 		Name:             fmt.Sprintf("connection_latency/nodes=%d/multiregion/password", numMultiRegionNodes),
 		Owner:            registry.OwnerSQLFoundations,
 		Benchmark:        true,
-		Cluster:          r.MakeClusterSpec(numMultiRegionNodes+loadNodes, spec.Geo(), spec.Zones(geoZonesStr)),
-		CompatibleClouds: registry.AllExceptAWS,
+		Cluster:          r.MakeClusterSpec(numMultiRegionNodes+loadNodes, spec.Geo(), spec.GCEZones(geoZonesStr)),
+		CompatibleClouds: registry.OnlyGCE,
 		Suites:           registry.Suites(registry.Nightly),
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
 			runConnectionLatencyTest(ctx, t, c, numMultiRegionNodes, numZones, true /*password*/)

pkg/cmd/roachtest/tests/decommissionbench.go

@@ -271,7 +271,7 @@ func registerDecommissionBenchSpec(r registry.Registry, benchSpec decommissionBe
 
 	if benchSpec.multiregion {
 		geoZones := []string{regionUsEast, regionUsWest, regionUsCentral}
-		specOptions = append(specOptions, spec.Zones(strings.Join(geoZones, ",")))
+		specOptions = append(specOptions, spec.GCEZones(strings.Join(geoZones, ",")))
 		specOptions = append(specOptions, spec.Geo())
 		extraNameParts = append(extraNameParts, "multi-region")
 	}

pkg/cmd/roachtest/tests/follower_reads.go

@@ -62,9 +62,9 @@ func registerFollowerReads(r registry.Registry) {
 				6, /* nodeCount */
 				spec.CPU(4),
 				spec.Geo(),
-				spec.Zones("us-east1-b,us-east1-b,us-east1-b,us-west1-b,us-west1-b,europe-west2-b"),
+				spec.GCEZones("us-east1-b,us-east1-b,us-east1-b,us-west1-b,us-west1-b,europe-west2-b"),
 			),
-			CompatibleClouds: registry.AllExceptAWS,
+			CompatibleClouds: registry.OnlyGCE,
 			Suites:           registry.Suites(registry.Nightly),
 			Leases:           registry.MetamorphicLeases,
 			Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {

pkg/cmd/roachtest/tests/import.go

@@ -195,8 +195,8 @@ func registerImportTPCC(r registry.Registry) {
 	r.Add(registry.TestSpec{
 		Name:              fmt.Sprintf("import/tpcc/warehouses=%d/geo", geoWarehouses),
 		Owner:             registry.OwnerSQLQueries,
-		Cluster:           r.MakeClusterSpec(8, spec.CPU(16), spec.Geo(), spec.Zones(geoZones)),
-		CompatibleClouds:  registry.AllExceptAWS,
+		Cluster:           r.MakeClusterSpec(8, spec.CPU(16), spec.Geo(), spec.GCEZones(geoZones)),
+		CompatibleClouds:  registry.OnlyGCE,
 		Suites:            registry.Suites(registry.Nightly),
 		Timeout:           5 * time.Hour,
 		EncryptionSupport: registry.EncryptionMetamorphic,

pkg/cmd/roachtest/tests/indexes.go

@@ -28,21 +28,28 @@ import (
 
 func registerNIndexes(r registry.Registry, secondaryIndexes int) {
 	const nodes = 6
-	geoZones := []string{"us-east1-b", "us-west1-b", "europe-west2-b"}
-	if r.MakeClusterSpec(1).Cloud == spec.AWS {
-		geoZones = []string{"us-east-2b", "us-west-1a", "eu-west-1a"}
-	}
-	geoZonesStr := strings.Join(geoZones, ",")
+	gceGeoZones := []string{"us-east1-b", "us-west1-b", "europe-west2-b"}
+	awsGeoZones := []string{"us-east-2b", "us-west-1a", "eu-west-1a"}
 	r.Add(registry.TestSpec{
-		Name:             fmt.Sprintf("indexes/%d/nodes=%d/multi-region", secondaryIndexes, nodes),
-		Owner:            registry.OwnerKV,
-		Benchmark:        true,
-		Cluster:          r.MakeClusterSpec(nodes+1, spec.CPU(16), spec.Geo(), spec.Zones(geoZonesStr)),
+		Name:      fmt.Sprintf("indexes/%d/nodes=%d/multi-region", secondaryIndexes, nodes),
+		Owner:     registry.OwnerKV,
+		Benchmark: true,
+		Cluster: r.MakeClusterSpec(
+			nodes+1,
+			spec.CPU(16),
+			spec.Geo(),
+			spec.GCEZones(strings.Join(gceGeoZones, ",")),
+			spec.AWSZones(strings.Join(awsGeoZones, ",")),
+		),
+		// TODO(radu): enable this test on AWS.
 		CompatibleClouds: registry.AllExceptAWS,
 		Suites:           registry.Suites(registry.Nightly),
 		// Uses CONFIGURE ZONE USING ... COPY FROM PARENT syntax.
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-			firstAZ := geoZones[0]
+			firstAZ := gceGeoZones[0]
+			if c.Spec().Cloud == spec.AWS {
+				firstAZ = awsGeoZones[0]
+			}
 			roachNodes := c.Range(1, nodes)
 			gatewayNodes := c.Range(1, nodes/3)
 			loadNode := c.Node(nodes + 1)

pkg/cmd/roachtest/tests/ledger.go

@@ -31,8 +31,8 @@ func registerLedger(r registry.Registry) {
 		Name:             fmt.Sprintf("ledger/nodes=%d/multi-az", nodes),
 		Owner:            registry.OwnerKV,
 		Benchmark:        true,
-		Cluster:          r.MakeClusterSpec(nodes+1, spec.CPU(16), spec.Geo(), spec.Zones(azs)),
-		CompatibleClouds: registry.AllExceptAWS,
+		Cluster:          r.MakeClusterSpec(nodes+1, spec.CPU(16), spec.Geo(), spec.GCEZones(azs)),
+		CompatibleClouds: registry.OnlyGCE,
 		Suites:           registry.Suites(registry.Nightly),
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
 			roachNodes := c.Range(1, nodes)

pkg/cmd/roachtest/tests/mixed_version_cdc.go

@@ -67,18 +67,13 @@ var (
 )
 
 func registerCDCMixedVersions(r registry.Registry) {
-	var zones string
-	if r.MakeClusterSpec(1).Cloud == spec.GCE {
-		// see rationale in definition of `teamcityAgentZone`
-		zones = teamcityAgentZone
-	}
 	r.Add(registry.TestSpec{
 		Name:  "cdc/mixed-versions",
 		Owner: registry.OwnerCDC,
 		// N.B. ARM64 is not yet supported, see https://github.com/cockroachdb/cockroach/issues/103888.
-		Cluster:          r.MakeClusterSpec(5, spec.Zones(zones), spec.Arch(vm.ArchAMD64)),
+		Cluster:          r.MakeClusterSpec(5, spec.GCEZones(teamcityAgentZone), spec.Arch(vm.ArchAMD64)),
 		Timeout:          60 * time.Minute,
-		CompatibleClouds: registry.AllExceptAWS,
+		CompatibleClouds: registry.OnlyGCE,
 		Suites:           registry.Suites(registry.Nightly),
 		RequiresLicense:  true,
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {

pkg/cmd/roachtest/tests/restore.go

@@ -529,7 +529,11 @@ func (hw hardwareSpecs) makeClusterSpecs(r registry.Registry, backupCloud string
 		addWorkloadNode++
 	}
 	if len(hw.zones) > 0 {
-		clusterOpts = append(clusterOpts, spec.Zones(strings.Join(hw.zones, ",")))
+		// Each test is set up to run on one specific cloud, so it's ok that the
+		// zones will only make sense for one of them.
+		// TODO(radu): clean this up.
+		clusterOpts = append(clusterOpts, spec.GCEZones(strings.Join(hw.zones, ",")))
+		clusterOpts = append(clusterOpts, spec.AWSZones(strings.Join(hw.zones, ",")))
 		clusterOpts = append(clusterOpts, spec.Geo())
 	}
 	if hw.ebsThroughput != 0 {

pkg/cmd/roachtest/tests/roachmart.go

@@ -74,8 +74,8 @@ func registerRoachmart(r registry.Registry) {
 		r.Add(registry.TestSpec{
 			Name:             fmt.Sprintf("roachmart/partition=%v", v),
 			Owner:            registry.OwnerKV,
-			Cluster:          r.MakeClusterSpec(9, spec.Geo(), spec.Zones("us-central1-b,us-west1-b,europe-west2-b")),
-			CompatibleClouds: registry.AllExceptAWS,
+			Cluster:          r.MakeClusterSpec(9, spec.Geo(), spec.GCEZones("us-central1-b,us-west1-b,europe-west2-b")),
+			CompatibleClouds: registry.OnlyGCE,
 			Suites:           registry.Suites(registry.Nightly),
 			Leases:           registry.MetamorphicLeases,
 			Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {

pkg/cmd/roachtest/tests/schemachange_random_load.go

@@ -15,7 +15,6 @@ import (
 	gosql "database/sql"
 	"fmt"
 	"path/filepath"
-	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
@@ -30,20 +29,17 @@ const (
 )
 
 func registerSchemaChangeRandomLoad(r registry.Registry) {
-	geoZones := []string{"us-east1-b", "us-west1-b", "europe-west2-b"}
-	if r.MakeClusterSpec(1).Cloud == spec.AWS {
-		geoZones = []string{"us-east-2b", "us-west-1a", "eu-west-1a"}
-	}
-	geoZonesStr := strings.Join(geoZones, ",")
 	r.Add(registry.TestSpec{
 		Name:      "schemachange/random-load",
 		Owner:     registry.OwnerSQLFoundations,
 		Benchmark: true,
 		Cluster: r.MakeClusterSpec(
 			3,
 			spec.Geo(),
-			spec.Zones(geoZonesStr),
+			spec.GCEZones("us-east1-b,us-west1-b,europe-west2-b"),
+			spec.AWSZones("us-east-2b,us-west-1a,eu-west-1a"),
 		),
+		// TODO(radu): enable this test on AWS.
 		CompatibleClouds: registry.AllExceptAWS,
 		Suites:           registry.Suites(registry.Nightly),
 		Leases:           registry.MetamorphicLeases,