From ca898ce33583144a024d51e42723cfc2c6ac1a2d Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Wed, 4 Oct 2023 13:11:10 -0700
Subject: [PATCH 1/5] roachtest: spec: simplify Option

This commit reduces boilerplate in the `Option` code by making it a
func instead of an interface with an apply func. This way each option
can define the function inline instead of having to define a type.

Epic: none
Release note: None
---
 pkg/cmd/roachtest/spec/cluster_spec.go |   4 +-
 pkg/cmd/roachtest/spec/option.go       | 189 +++++++++----------------
 2 files changed, 68 insertions(+), 125 deletions(-)

diff --git a/pkg/cmd/roachtest/spec/cluster_spec.go b/pkg/cmd/roachtest/spec/cluster_spec.go
index 739523b1fa19..c49418ff8bd8 100644
--- a/pkg/cmd/roachtest/spec/cluster_spec.go
+++ b/pkg/cmd/roachtest/spec/cluster_spec.go
@@ -101,9 +101,9 @@ type ClusterSpec struct {
 // MakeClusterSpec makes a ClusterSpec.
 func MakeClusterSpec(cloud string, instanceType string, nodeCount int, opts ...Option) ClusterSpec {
 	spec := ClusterSpec{Cloud: cloud, InstanceType: instanceType, NodeCount: nodeCount}
-	defaultOpts := []Option{CPU(4), nodeLifetimeOption(12 * time.Hour), ReuseAny()}
+	defaultOpts := []Option{CPU(4), nodeLifetime(12 * time.Hour), ReuseAny()}
 	for _, o := range append(defaultOpts, opts...) {
-		o.apply(&spec)
+		o(&spec)
 	}
 	return spec
 }
diff --git a/pkg/cmd/roachtest/spec/option.go b/pkg/cmd/roachtest/spec/option.go
index 59918d708a3e..c903e38401c6 100644
--- a/pkg/cmd/roachtest/spec/option.go
+++ b/pkg/cmd/roachtest/spec/option.go
@@ -16,127 +16,90 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachprod/vm"
 )
 
-// Option is the interface satisfied by options to MakeClusterSpec.
-type Option interface {
-	apply(spec *ClusterSpec)
-}
-
-type cloudOption string
-
-func (o cloudOption) apply(spec *ClusterSpec) {
-	spec.Cloud = string(o)
-}
+// Option for MakeClusterSpec.
+type Option func(spec *ClusterSpec)
 
 // Cloud controls what cloud is used to create the cluster.
 func Cloud(s string) Option {
-	return cloudOption(s)
-}
-
-type archOption string
-
-func (o archOption) apply(spec *ClusterSpec) {
-	spec.Arch = vm.CPUArch(o)
+	return func(spec *ClusterSpec) {
+		spec.Cloud = s
+	}
 }
 
-// Request specific CPU architecture.
+// Arch requests a specific CPU architecture.
+//
+// Note that it is not guaranteed that this architecture will be used (e.g. if
+// the requested machine size isn't available in this architecture).
+//
+// TODO(radu): add a flag to indicate whether it's a preference or a requirement.
 func Arch(arch vm.CPUArch) Option {
-	return archOption(arch)
-}
-
-type nodeCPUOption int
-
-func (o nodeCPUOption) apply(spec *ClusterSpec) {
-	spec.CPUs = int(o)
+	return func(spec *ClusterSpec) {
+		spec.Arch = arch
+	}
 }
 
-// CPU is a node option which requests nodes with the specified number of CPUs.
+// CPU sets the number of CPUs for each node.
 func CPU(n int) Option {
-	return nodeCPUOption(n)
-}
-
-type nodeMemOption MemPerCPU
-
-func (o nodeMemOption) apply(spec *ClusterSpec) {
-	spec.Mem = MemPerCPU(o)
+	return func(spec *ClusterSpec) {
+		spec.CPUs = n
+	}
 }
 
 // Mem requests nodes with low/standard/high ratio of memory per CPU.
 func Mem(level MemPerCPU) Option {
-	return nodeMemOption(level)
-}
-
-type volumeSizeOption int
-
-func (o volumeSizeOption) apply(spec *ClusterSpec) {
-	spec.VolumeSize = int(o)
+	return func(spec *ClusterSpec) {
+		spec.Mem = level
+	}
 }
 
 // VolumeSize is the size in GB of the disk volume.
 func VolumeSize(n int) Option {
-	return volumeSizeOption(n)
-}
-
-type nodeSSDOption int
-
-func (o nodeSSDOption) apply(spec *ClusterSpec) {
-	spec.SSDs = int(o)
+	return func(spec *ClusterSpec) {
+		spec.VolumeSize = n
+	}
 }
 
 // SSD is a node option which requests nodes with the specified number of SSDs.
 func SSD(n int) Option {
-	return nodeSSDOption(n)
-}
-
-type raid0Option bool
-
-func (o raid0Option) apply(spec *ClusterSpec) {
-	spec.RAID0 = bool(o)
+	return func(spec *ClusterSpec) {
+		spec.SSDs = n
+	}
 }
 
 // RAID0 enables RAID 0 striping across all disks on the node.
 func RAID0(enabled bool) Option {
-	return raid0Option(enabled)
+	return func(spec *ClusterSpec) {
+		spec.RAID0 = enabled
+	}
 }
 
-type nodeGeoOption struct{}
-
-func (o nodeGeoOption) apply(spec *ClusterSpec) {
-	spec.Geo = true
-}
-
-// Geo is a node option which requests Geo-distributed nodes.
+// Geo requests Geo-distributed nodes.
 func Geo() Option {
-	return nodeGeoOption{}
-}
-
-type nodeZonesOption string
-
-func (o nodeZonesOption) apply(spec *ClusterSpec) {
-	spec.Zones = string(o)
+	return func(spec *ClusterSpec) {
+		spec.Geo = true
+	}
 }
 
 // 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(s string) Option {
-	return nodeZonesOption(s)
+func Zones(zones string) Option {
+	return func(spec *ClusterSpec) {
+		spec.Zones = zones
+	}
 }
 
-type nodeLifetimeOption time.Duration
-
-func (o nodeLifetimeOption) apply(spec *ClusterSpec) {
-	spec.Lifetime = time.Duration(o)
-}
-
-type gatherCoresOption struct{}
-
-func (o gatherCoresOption) apply(spec *ClusterSpec) {
-	spec.GatherCores = true
+func nodeLifetime(lifetime time.Duration) Option {
+	return func(spec *ClusterSpec) {
+		spec.Lifetime = lifetime
+	}
 }
 
 // GatherCores enables core gathering after test runs.
 func GatherCores() Option {
-	return gatherCoresOption{}
+	return func(spec *ClusterSpec) {
+		spec.GatherCores = true
+	}
 }
 
 // clusterReusePolicy indicates what clusters a particular test can run on and
@@ -182,69 +145,47 @@ func (ReusePolicyAny) clusterReusePolicy()    {}
 func (ReusePolicyNone) clusterReusePolicy()   {}
 func (ReusePolicyTagged) clusterReusePolicy() {}
 
-type clusterReusePolicyOption struct {
-	p clusterReusePolicy
-}
-
 // ReuseAny is an Option that specifies a cluster with ReusePolicyAny.
 func ReuseAny() Option {
-	return clusterReusePolicyOption{p: ReusePolicyAny{}}
+	return func(spec *ClusterSpec) {
+		spec.ReusePolicy = ReusePolicyAny{}
+	}
 }
 
 // ReuseNone is an Option that specifies a cluster with ReusePolicyNone.
 func ReuseNone() Option {
-	return clusterReusePolicyOption{p: ReusePolicyNone{}}
+	return func(spec *ClusterSpec) {
+		spec.ReusePolicy = ReusePolicyNone{}
+	}
 }
 
 // ReuseTagged is an Option that specifies a cluster with ReusePolicyTagged.
 func ReuseTagged(tag string) Option {
-	return clusterReusePolicyOption{p: ReusePolicyTagged{Tag: tag}}
-}
-
-func (p clusterReusePolicyOption) apply(spec *ClusterSpec) {
-	spec.ReusePolicy = p.p
-}
-
-type preferLocalSSDOption bool
-
-func (o preferLocalSSDOption) apply(spec *ClusterSpec) {
-	spec.PreferLocalSSD = bool(o)
+	return func(spec *ClusterSpec) {
+		spec.ReusePolicy = ReusePolicyTagged{Tag: tag}
+	}
 }
 
 // PreferLocalSSD specifies whether to prefer using local SSD, when possible.
 func PreferLocalSSD(prefer bool) Option {
-	return preferLocalSSDOption(prefer)
-}
-
-type terminateOnMigrationOption struct{}
-
-func (o terminateOnMigrationOption) apply(spec *ClusterSpec) {
-	spec.TerminateOnMigration = true
+	return func(spec *ClusterSpec) {
+		spec.PreferLocalSSD = prefer
+	}
 }
 
 // TerminateOnMigration ensures VM is terminated in case GCE triggers a live migration.
 func TerminateOnMigration() Option {
-	return &terminateOnMigrationOption{}
-}
-
-type setFileSystem struct {
-	fs fileSystemType
-}
-
-func (s *setFileSystem) apply(spec *ClusterSpec) {
-	spec.FileSystem = s.fs
+	return func(spec *ClusterSpec) {
+		spec.TerminateOnMigration = true
+	}
 }
 
 // SetFileSystem is an Option which can be used to set
 // the underlying file system to be used.
 func SetFileSystem(fs fileSystemType) Option {
-	return &setFileSystem{fs}
-}
-
-type randomlyUseZfs struct{}
-
-func (r *randomlyUseZfs) apply(spec *ClusterSpec) {
-	spec.RandomlyUseZfs = true
+	return func(spec *ClusterSpec) {
+		spec.FileSystem = fs
+	}
 }
 
 // RandomlyUseZfs is an Option which randomly picks
@@ -252,7 +193,9 @@ func (r *randomlyUseZfs) apply(spec *ClusterSpec) {
 // about 20% of the time.
 // Zfs is only picked if the cloud is gce.
 func RandomlyUseZfs() Option {
-	return &randomlyUseZfs{}
+	return func(spec *ClusterSpec) {
+		spec.RandomlyUseZfs = true
+	}
 }
 
 type ubuntuVersion vm.UbuntuVersion

From 1ef46729e78fc0b28534d8820c4bedc66528f96d Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Wed, 4 Oct 2023 13:13:11 -0700
Subject: [PATCH 2/5] roachtest: spec: provide options for GCE/AWS settings

This commit moves GCE and AWS specific settings to their own
inline structs and adds Options for them.

Epic: none
Release note: None
---
 pkg/cmd/roachtest/spec/cluster_spec.go | 19 +++++++++++------
 pkg/cmd/roachtest/spec/option.go       | 28 +++++++++++++++++++++-----
 pkg/cmd/roachtest/tests/restore.go     |  5 ++---
 3 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/pkg/cmd/roachtest/spec/cluster_spec.go b/pkg/cmd/roachtest/spec/cluster_spec.go
index c49418ff8bd8..5e46152bff9b 100644
--- a/pkg/cmd/roachtest/spec/cluster_spec.go
+++ b/pkg/cmd/roachtest/spec/cluster_spec.go
@@ -92,10 +92,17 @@ type ClusterSpec struct {
 
 	GatherCores bool
 
-	// AWS-specific arguments.
-	//
-	// AWSVolumeThroughput is the min provisioned EBS volume throughput.
-	AWSVolumeThroughput int
+	// GCE-specific arguments. These values apply only on clusters instantiated on GCE.
+	GCE struct {
+		MinCPUPlatform string
+		VolumeType     string
+	}
+
+	// AWS-specific arguments. These values apply only on clusters instantiated on AWS.
+	AWS struct {
+		// VolumeThroughput is the min provisioned EBS volume throughput.
+		VolumeThroughput int
+	}
 }
 
 // MakeClusterSpec makes a ClusterSpec.
@@ -321,12 +328,12 @@ func (s *ClusterSpec) RoachprodOpts(
 	var providerOpts vm.ProviderOpts
 	switch s.Cloud {
 	case AWS:
-		providerOpts = getAWSOpts(machineType, zones, s.VolumeSize, s.AWSVolumeThroughput,
+		providerOpts = getAWSOpts(machineType, zones, s.VolumeSize, s.AWS.VolumeThroughput,
 			createVMOpts.SSDOpts.UseLocalSSD)
 	case GCE:
 		providerOpts = getGCEOpts(machineType, zones, s.VolumeSize, ssdCount,
 			createVMOpts.SSDOpts.UseLocalSSD, s.RAID0, s.TerminateOnMigration,
-			"" /* minCPUPlatform */, vm.ParseArch(createVMOpts.Arch),
+			s.GCE.MinCPUPlatform, vm.ParseArch(createVMOpts.Arch),
 		)
 	case Azure:
 		providerOpts = getAzureOpts(machineType, zones)
diff --git a/pkg/cmd/roachtest/spec/option.go b/pkg/cmd/roachtest/spec/option.go
index c903e38401c6..35b70ad40fc0 100644
--- a/pkg/cmd/roachtest/spec/option.go
+++ b/pkg/cmd/roachtest/spec/option.go
@@ -198,13 +198,31 @@ func RandomlyUseZfs() Option {
 	}
 }
 
-type ubuntuVersion vm.UbuntuVersion
+// GCEMinCPUPlatform sets the minimum CPU platform when the cluster is on GCE.
+func GCEMinCPUPlatform(platform string) Option {
+	return func(spec *ClusterSpec) {
+		spec.GCE.MinCPUPlatform = platform
+	}
+}
 
-func (u ubuntuVersion) apply(spec *ClusterSpec) {
-	spec.UbuntuVersion = vm.UbuntuVersion(u)
+// GCEVolumeType sets the volume type when the cluster is on GCE.
+func GCEVolumeType(volumeType string) Option {
+	return func(spec *ClusterSpec) {
+		spec.GCE.VolumeType = volumeType
+	}
+}
+
+// AWSVolumeThroughput sets the minimum provisioned EBS volume throughput when
+// the cluster is on AWS.
+func AWSVolumeThroughput(throughput int) Option {
+	return func(spec *ClusterSpec) {
+		spec.AWS.VolumeThroughput = throughput
+	}
 }
 
 // UbuntuVersion controls what Ubuntu Version is used to create the cluster.
-func UbuntuVersion(u vm.UbuntuVersion) Option {
-	return ubuntuVersion(u)
+func UbuntuVersion(version vm.UbuntuVersion) Option {
+	return func(spec *ClusterSpec) {
+		spec.UbuntuVersion = version
+	}
 }
diff --git a/pkg/cmd/roachtest/tests/restore.go b/pkg/cmd/roachtest/tests/restore.go
index 835b392d5e65..01a89153b8ea 100644
--- a/pkg/cmd/roachtest/tests/restore.go
+++ b/pkg/cmd/roachtest/tests/restore.go
@@ -484,11 +484,10 @@ func (hw hardwareSpecs) makeClusterSpecs(r registry.Registry, backupCloud string
 		clusterOpts = append(clusterOpts, spec.Zones(strings.Join(hw.zones, ",")))
 		clusterOpts = append(clusterOpts, spec.Geo())
 	}
-	s := r.MakeClusterSpec(hw.nodes, clusterOpts...)
-
 	if hw.ebsThroughput != 0 {
-		s.AWSVolumeThroughput = hw.ebsThroughput
+		clusterOpts = append(clusterOpts, spec.AWSVolumeThroughput(hw.ebsThroughput))
 	}
+	s := r.MakeClusterSpec(hw.nodes, clusterOpts...)
 
 	if backupCloud == spec.AWS && s.Cloud == spec.AWS && s.VolumeSize != 0 {
 		// Work around an issue that RAID0s local NVMe and GP3 storage together:

From 4fdbeec588b22b0813a05beec77eda137cae2c88 Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Wed, 4 Oct 2023 13:38:33 -0700
Subject: [PATCH 3/5] roachtest: clean up instance type specification

`ClusterSpec` now provides options for GCE and AWS machine types.

Epic: none
Release note: None
---
 pkg/cmd/roachtest/cluster_test.go       |  8 ++--
 pkg/cmd/roachtest/spec/cluster_spec.go  | 50 ++++++++++++++++---------
 pkg/cmd/roachtest/spec/machine_type.go  | 40 +++++++++++---------
 pkg/cmd/roachtest/spec/option.go        | 14 +++++++
 pkg/cmd/roachtest/test_registry_test.go |  3 --
 pkg/cmd/roachtest/tests/restore.go      |  6 +--
 6 files changed, 76 insertions(+), 45 deletions(-)

diff --git a/pkg/cmd/roachtest/cluster_test.go b/pkg/cmd/roachtest/cluster_test.go
index ca80d8ec9ac7..747141753818 100644
--- a/pkg/cmd/roachtest/cluster_test.go
+++ b/pkg/cmd/roachtest/cluster_test.go
@@ -320,15 +320,15 @@ func TestAWSMachineTypeNew(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(fmt.Sprintf("%d/%s/%t/%s", tc.cpus, tc.mem, tc.localSSD, tc.arch), func(t *testing.T) {
-			machineType, selectedArch := spec.AWSMachineTypeNew(tc.cpus, tc.mem, tc.localSSD, tc.arch)
+			machineType, selectedArch := spec.SelectAWSMachineTypeNew(tc.cpus, tc.mem, tc.localSSD, tc.arch)
 
 			require.Equal(t, tc.expectedMachineType, machineType)
 			require.Equal(t, tc.expectedArch, selectedArch)
 		})
 	}
 	// spec.Low is not supported.
-	require.Panics(t, func() { spec.AWSMachineTypeNew(4, spec.Low, false, vm.ArchAMD64) })
-	require.Panics(t, func() { spec.AWSMachineTypeNew(16, spec.Low, false, vm.ArchARM64) })
+	require.Panics(t, func() { spec.SelectAWSMachineTypeNew(4, spec.Low, false, vm.ArchAMD64) })
+	require.Panics(t, func() { spec.SelectAWSMachineTypeNew(16, spec.Low, false, vm.ArchARM64) })
 }
 
 // TODO(srosenberg): restore the change in https://github.com/cockroachdb/cockroach/pull/111140 after 23.2 branch cut.
@@ -415,7 +415,7 @@ func TestGCEMachineTypeNew(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(fmt.Sprintf("%d/%s/%s", tc.cpus, tc.mem, tc.arch), func(t *testing.T) {
-			machineType, selectedArch := spec.GCEMachineTypeNew(tc.cpus, tc.mem, tc.arch)
+			machineType, selectedArch := spec.SelectGCEMachineTypeNew(tc.cpus, tc.mem, tc.arch)
 
 			require.Equal(t, tc.expectedMachineType, machineType)
 			require.Equal(t, tc.expectedArch, selectedArch)
diff --git a/pkg/cmd/roachtest/spec/cluster_spec.go b/pkg/cmd/roachtest/spec/cluster_spec.go
index 5e46152bff9b..65aabf67aa9f 100644
--- a/pkg/cmd/roachtest/spec/cluster_spec.go
+++ b/pkg/cmd/roachtest/spec/cluster_spec.go
@@ -65,11 +65,14 @@ type ClusterSpec struct {
 	// TODO(#104029): We should remove the Cloud field; the tests now specify
 	// their compatible clouds.
 	Cloud string
-	Arch  vm.CPUArch // CPU architecture; auto-chosen if left empty
-	// TODO(radu): An InstanceType can only make sense in the context of a
-	// specific cloud. We should replace this with cloud-specific arguments.
-	InstanceType string // auto-chosen if left empty
-	NodeCount    int
+
+	// TODO(radu): defaultInstanceType is the default machine type used (unless
+	// overridden by GCE.MachineType or AWS.MachineType); it does not belong in
+	// the spec.
+	defaultInstanceType string
+
+	Arch      vm.CPUArch // CPU architecture; auto-chosen if left empty
+	NodeCount int
 	// CPUs is the number of CPUs per node.
 	CPUs                 int
 	Mem                  MemPerCPU
@@ -94,12 +97,14 @@ type ClusterSpec struct {
 
 	// GCE-specific arguments. These values apply only on clusters instantiated on GCE.
 	GCE struct {
+		MachineType    string
 		MinCPUPlatform string
 		VolumeType     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
 	}
@@ -107,7 +112,7 @@ type ClusterSpec struct {
 
 // MakeClusterSpec makes a ClusterSpec.
 func MakeClusterSpec(cloud string, instanceType string, nodeCount int, opts ...Option) ClusterSpec {
-	spec := ClusterSpec{Cloud: cloud, InstanceType: instanceType, NodeCount: nodeCount}
+	spec := ClusterSpec{Cloud: cloud, defaultInstanceType: instanceType, NodeCount: nodeCount}
 	defaultOpts := []Option{CPU(4), nodeLifetime(12 * time.Hour), ReuseAny()}
 	for _, o := range append(defaultOpts, opts...) {
 		o(&spec)
@@ -254,31 +259,42 @@ func (s *ClusterSpec) RoachprodOpts(
 
 	createVMOpts.GeoDistributed = s.Geo
 	createVMOpts.Arch = string(arch)
-	machineType := s.InstanceType
 	ssdCount := s.SSDs
 
+	machineType := s.defaultInstanceType
+	switch s.Cloud {
+	case AWS:
+		if s.AWS.MachineType != "" {
+			machineType = s.AWS.MachineType
+		}
+	case GCE:
+		if s.GCE.MachineType != "" {
+			machineType = s.GCE.MachineType
+		}
+	}
+
 	if s.CPUs != 0 {
 		// Default to the user-supplied machine type, if any.
 		// Otherwise, pick based on requested CPU count.
 		var selectedArch vm.CPUArch
 
-		if len(machineType) == 0 {
+		if machineType == "" {
 			// If no machine type was specified, choose one
 			// based on the cloud and CPU count.
 			switch s.Cloud {
 			case AWS:
-				machineType, selectedArch = AWSMachineType(s.CPUs, s.Mem, s.PreferLocalSSD && s.VolumeSize == 0, arch)
+				machineType, selectedArch = SelectAWSMachineType(s.CPUs, s.Mem, s.PreferLocalSSD && s.VolumeSize == 0, arch)
 			case GCE:
-				machineType, selectedArch = GCEMachineType(s.CPUs, s.Mem, arch)
+				machineType, selectedArch = SelectGCEMachineType(s.CPUs, s.Mem, arch)
 			case Azure:
-				machineType = AzureMachineType(s.CPUs, s.Mem)
+				machineType = SelectAzureMachineType(s.CPUs, s.Mem)
+			}
+			if selectedArch != "" && selectedArch != arch {
+				// TODO(srosenberg): we need a better way to monitor the rate of this mismatch, i.e.,
+				// other than grepping cluster creation logs.
+				fmt.Printf("WARN: requested arch %s for machineType %s, but selected %s\n", arch, machineType, selectedArch)
+				createVMOpts.Arch = string(selectedArch)
 			}
-		}
-		if selectedArch != "" && selectedArch != arch {
-			// TODO(srosenberg): we need a better way to monitor the rate of this mismatch, i.e.,
-			// other than grepping cluster creation logs.
-			fmt.Printf("WARN: requested arch %s for machineType %s, but selected %s\n", arch, machineType, selectedArch)
-			createVMOpts.Arch = string(selectedArch)
 		}
 
 		// Local SSD can only be requested
diff --git a/pkg/cmd/roachtest/spec/machine_type.go b/pkg/cmd/roachtest/spec/machine_type.go
index 0f3400c827e6..15e76f04f784 100644
--- a/pkg/cmd/roachtest/spec/machine_type.go
+++ b/pkg/cmd/roachtest/spec/machine_type.go
@@ -17,20 +17,21 @@ import (
 )
 
 // TODO(srosenberg): restore the change in https://github.com/cockroachdb/cockroach/pull/111140 after 23.2 branch cut.
-func AWSMachineType(
+func SelectAWSMachineType(
 	cpus int, mem MemPerCPU, shouldSupportLocalSSD bool, arch vm.CPUArch,
 ) (string, vm.CPUArch) {
-	return AWSMachineTypeOld(cpus, mem, arch)
+	return SelectAWSMachineTypeOld(cpus, mem, arch)
 }
 
 // TODO(srosenberg): restore the change in https://github.com/cockroachdb/cockroach/pull/111140 after 23.2 branch cut.
-func GCEMachineType(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPUArch) {
-	return GCEMachineTypeOld(cpus, mem, arch)
+func SelectGCEMachineType(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPUArch) {
+	return SelectGCEMachineTypeOld(cpus, mem, arch)
 }
 
-// AWSMachineType selects a machine type given the desired number of CPUs and
-// memory per CPU ratio. Also returns the architecture of the selected machine type.
-func AWSMachineTypeOld(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPUArch) {
+// SelectAWSMachineType selects a machine type given the desired number of CPUs
+// and memory per CPU ratio. Also returns the architecture of the selected
+// machine type.
+func SelectAWSMachineTypeOld(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPUArch) {
 	// TODO(erikgrinaker): These have significantly less RAM than
 	// their GCE counterparts. Consider harmonizing them.
 	family := "c6id" // 2 GB RAM per CPU
@@ -87,8 +88,9 @@ func AWSMachineTypeOld(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPU
 	return fmt.Sprintf("%s.%s", family, size), selectedArch
 }
 
-// AWSMachineType selects a machine type given the desired number of CPUs, memory per CPU,
-// support for locally-attached SSDs and CPU architecture. It returns a compatible machine type and its architecture.
+// SelectAWSMachineType selects a machine type given the desired number of CPUs,
+// memory per CPU, support for locally-attached SSDs and CPU architecture. It
+// returns a compatible machine type and its architecture.
 //
 // When MemPerCPU is Standard, the memory per CPU ratio is 4 GB. For High, it is 8 GB.
 // For Auto, it's 4 GB up to and including 16 CPUs, then 2 GB. Low is not supported.
@@ -99,7 +101,7 @@ func AWSMachineTypeOld(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPU
 //
 // At the time of writing, the intel machines are all third-generation Xeon, "Ice Lake" which are isomorphic to
 // GCE's n2-(standard|highmem|custom) _with_ --minimum-cpu-platform="Intel Ice Lake" (roachprod's default).
-func AWSMachineTypeNew(
+func SelectAWSMachineTypeNew(
 	cpus int, mem MemPerCPU, shouldSupportLocalSSD bool, arch vm.CPUArch,
 ) (string, vm.CPUArch) {
 	family := "m6i" // 4 GB RAM per CPU
@@ -176,9 +178,10 @@ func AWSMachineTypeNew(
 	return fmt.Sprintf("%s.%s", family, size), selectedArch
 }
 
-// GCEMachineType selects a machine type given the desired number of CPUs and
-// memory per CPU ratio. Also returns the architecture of the selected machine type.
-func GCEMachineTypeOld(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPUArch) {
+// SelectGCEMachineType selects a machine type given the desired number of CPUs
+// and memory per CPU ratio. Also returns the architecture of the selected
+// machine type.
+func SelectGCEMachineTypeOld(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPUArch) {
 	// TODO(peter): This is awkward: at or below 16 cpus, use n2-standard so that
 	// the machines have a decent amount of RAM. We could use custom machine
 	// configurations, but the rules for the amount of RAM per CPU need to be
@@ -215,8 +218,9 @@ func GCEMachineTypeOld(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPU
 	return fmt.Sprintf("%s-%s-%d", series, kind, cpus), selectedArch
 }
 
-// GCEMachineType selects a machine type given the desired number of CPUs, memory per CPU, and CPU architecture.
-// It returns a compatible machine type and its architecture.
+// SelectGCEMachineType selects a machine type given the desired number of CPUs,
+// memory per CPU, and CPU architecture.  It returns a compatible machine type
+// and its architecture.
 //
 // When MemPerCPU is Standard, the memory per CPU ratio is 4 GB. For High, it is 8 GB.
 // For Auto, it's 4 GB up to and including 16 CPUs, then 2 GB. Low is 1 GB.
@@ -228,7 +232,7 @@ func GCEMachineTypeOld(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPU
 // At the time of writing, the intel machines are all third-generation xeon, "Ice Lake" assuming
 // --minimum-cpu-platform="Intel Ice Lake" (roachprod's default). This is isomorphic to AWS's m6i or c6i.
 // The only exception is low memory machines (n2-highcpu-xxx), which aren't available in AWS.
-func GCEMachineTypeNew(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPUArch) {
+func SelectGCEMachineTypeNew(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPUArch) {
 	series := "n2"
 	selectedArch := vm.ArchAMD64
 
@@ -286,9 +290,9 @@ func GCEMachineTypeNew(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPU
 	return fmt.Sprintf("%s-%s-%d", series, kind, cpus), selectedArch
 }
 
-// AzureMachineType selects a machine type given the desired number of CPUs and
+// SelectAzureMachineType selects a machine type given the desired number of CPUs and
 // memory per CPU ratio.
-func AzureMachineType(cpus int, mem MemPerCPU) string {
+func SelectAzureMachineType(cpus int, mem MemPerCPU) string {
 	if mem != Auto && mem != Standard {
 		panic(fmt.Sprintf("custom memory per CPU not implemented for Azure, memory ratio requested: %d", mem))
 	}
diff --git a/pkg/cmd/roachtest/spec/option.go b/pkg/cmd/roachtest/spec/option.go
index 35b70ad40fc0..0101aecf7d7e 100644
--- a/pkg/cmd/roachtest/spec/option.go
+++ b/pkg/cmd/roachtest/spec/option.go
@@ -198,6 +198,13 @@ func RandomlyUseZfs() Option {
 	}
 }
 
+// GCEMachineType sets the machine (instance) type when the cluster is on GCE.
+func GCEMachineType(machineType string) Option {
+	return func(spec *ClusterSpec) {
+		spec.GCE.MachineType = machineType
+	}
+}
+
 // GCEMinCPUPlatform sets the minimum CPU platform when the cluster is on GCE.
 func GCEMinCPUPlatform(platform string) Option {
 	return func(spec *ClusterSpec) {
@@ -212,6 +219,13 @@ func GCEVolumeType(volumeType string) Option {
 	}
 }
 
+// AWSMachineType sets the machine (instance) type when the cluster is on AWS.
+func AWSMachineType(machineType string) Option {
+	return func(spec *ClusterSpec) {
+		spec.AWS.MachineType = machineType
+	}
+}
+
 // AWSVolumeThroughput sets the minimum provisioned EBS volume throughput when
 // the cluster is on AWS.
 func AWSVolumeThroughput(throughput int) Option {
diff --git a/pkg/cmd/roachtest/test_registry_test.go b/pkg/cmd/roachtest/test_registry_test.go
index bf19ef33c6db..aa42782fe79f 100644
--- a/pkg/cmd/roachtest/test_registry_test.go
+++ b/pkg/cmd/roachtest/test_registry_test.go
@@ -32,7 +32,6 @@ func TestMakeTestRegistry(t *testing.T) {
 		s := r.MakeClusterSpec(100, spec.Geo(), spec.Zones("zone99"), spec.CPU(12),
 			spec.PreferLocalSSD(true))
 		require.EqualValues(t, 100, s.NodeCount)
-		require.Equal(t, "foo", s.InstanceType)
 		require.True(t, s.Geo)
 		require.Equal(t, "zone99", s.Zones)
 		require.EqualValues(t, 12, s.CPUs)
@@ -40,13 +39,11 @@ func TestMakeTestRegistry(t *testing.T) {
 
 		s = r.MakeClusterSpec(100, spec.CPU(4), spec.TerminateOnMigration())
 		require.EqualValues(t, 100, s.NodeCount)
-		require.Equal(t, "foo", s.InstanceType)
 		require.EqualValues(t, 4, s.CPUs)
 		require.True(t, s.TerminateOnMigration)
 
 		s = r.MakeClusterSpec(10, spec.CPU(16), spec.Arch(vm.ArchARM64))
 		require.EqualValues(t, 10, s.NodeCount)
-		require.Equal(t, "foo", s.InstanceType)
 		require.EqualValues(t, 16, s.CPUs)
 		require.EqualValues(t, vm.ArchARM64, s.Arch)
 	})
diff --git a/pkg/cmd/roachtest/tests/restore.go b/pkg/cmd/roachtest/tests/restore.go
index 01a89153b8ea..8cff79545767 100644
--- a/pkg/cmd/roachtest/tests/restore.go
+++ b/pkg/cmd/roachtest/tests/restore.go
@@ -489,13 +489,13 @@ func (hw hardwareSpecs) makeClusterSpecs(r registry.Registry, backupCloud string
 	}
 	s := r.MakeClusterSpec(hw.nodes, clusterOpts...)
 
-	if backupCloud == spec.AWS && s.Cloud == spec.AWS && s.VolumeSize != 0 {
+	if backupCloud == spec.AWS && s.VolumeSize != 0 {
 		// Work around an issue that RAID0s local NVMe and GP3 storage together:
 		// https://github.com/cockroachdb/cockroach/issues/98783.
 		//
 		// TODO(srosenberg): Remove this workaround when 98783 is addressed.
-		s.InstanceType, _ = spec.AWSMachineType(s.CPUs, s.Mem, s.PreferLocalSSD && s.VolumeSize == 0, vm.ArchAMD64)
-		s.InstanceType = strings.Replace(s.InstanceType, "d.", ".", 1)
+		s.AWS.MachineType, _ = spec.SelectAWSMachineType(s.CPUs, s.Mem, s.PreferLocalSSD && s.VolumeSize == 0, vm.ArchAMD64)
+		s.AWS.MachineType = strings.Replace(s.AWS.MachineType, "d.", ".", 1)
 		s.Arch = vm.ArchAMD64
 	}
 	return s

From 75632e3a7ac7cb53a286682c9369d7d40becd28a Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Wed, 4 Oct 2023 13:38:35 -0700
Subject: [PATCH 4/5] 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        | 23 ++++++++++++--
 pkg/cmd/roachtest/spec/option.go              | 30 +++++++++++++++----
 pkg/cmd/roachtest/test_registry.go            |  2 +-
 pkg/cmd/roachtest/test_registry_test.go       |  3 +-
 pkg/cmd/roachtest/tests/connection_latency.go | 12 ++++----
 pkg/cmd/roachtest/tests/decommissionbench.go  |  2 +-
 pkg/cmd/roachtest/tests/follower_reads.go     |  4 +--
 pkg/cmd/roachtest/tests/import.go             |  4 +--
 pkg/cmd/roachtest/tests/indexes.go            | 27 ++++++++++-------
 pkg/cmd/roachtest/tests/ledger.go             |  4 +--
 pkg/cmd/roachtest/tests/mixed_version_cdc.go  |  9 ++----
 pkg/cmd/roachtest/tests/restore.go            |  6 +++-
 pkg/cmd/roachtest/tests/roachmart.go          |  4 +--
 .../tests/schemachange_random_load.go         |  9 ++----
 pkg/cmd/roachtest/tests/tpcc.go               | 12 ++++----
 15 files changed, 95 insertions(+), 56 deletions(-)

diff --git a/pkg/cmd/roachtest/spec/cluster_spec.go b/pkg/cmd/roachtest/spec/cluster_spec.go
index 65aabf67aa9f..b3b05032a72a 100644
--- a/pkg/cmd/roachtest/spec/cluster_spec.go
+++ b/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
@@ -100,6 +103,7 @@ type ClusterSpec struct {
 		MachineType    string
 		MinCPUPlatform string
 		VolumeType     string
+		Zones          string
 	}
 
 	// AWS-specific arguments. These values apply only on clusters instantiated on AWS.
@@ -107,6 +111,7 @@ type ClusterSpec struct {
 		MachineType string
 		// VolumeThroughput is the min provisioned EBS volume throughput.
 		VolumeThroughput int
+		Zones            string
 	}
 }
 
@@ -328,9 +333,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]
 		}
diff --git a/pkg/cmd/roachtest/spec/option.go b/pkg/cmd/roachtest/spec/option.go
index 0101aecf7d7e..da69b07be546 100644
--- a/pkg/cmd/roachtest/spec/option.go
+++ b/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) {
@@ -234,6 +243,17 @@ func AWSVolumeThroughput(throughput int) Option {
 	}
 }
 
+// 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
+	}
+}
+
 // UbuntuVersion controls what Ubuntu Version is used to create the cluster.
 func UbuntuVersion(version vm.UbuntuVersion) Option {
 	return func(spec *ClusterSpec) {
diff --git a/pkg/cmd/roachtest/test_registry.go b/pkg/cmd/roachtest/test_registry.go
index 6c704f526005..60dab1e3c32a 100644
--- a/pkg/cmd/roachtest/test_registry.go
+++ b/pkg/cmd/roachtest/test_registry.go
@@ -89,7 +89,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...)
diff --git a/pkg/cmd/roachtest/test_registry_test.go b/pkg/cmd/roachtest/test_registry_test.go
index aa42782fe79f..f78cd24ac72f 100644
--- a/pkg/cmd/roachtest/test_registry_test.go
+++ b/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)
 
diff --git a/pkg/cmd/roachtest/tests/connection_latency.go b/pkg/cmd/roachtest/tests/connection_latency.go
index 2d1170d54936..3af52b96a12c 100644
--- a/pkg/cmd/roachtest/tests/connection_latency.go
+++ b/pkg/cmd/roachtest/tests/connection_latency.go
@@ -119,8 +119,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*/)
@@ -137,8 +137,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*/)
@@ -149,8 +149,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*/)
diff --git a/pkg/cmd/roachtest/tests/decommissionbench.go b/pkg/cmd/roachtest/tests/decommissionbench.go
index d99ce3dfc11b..82f3c43bf6e1 100644
--- a/pkg/cmd/roachtest/tests/decommissionbench.go
+++ b/pkg/cmd/roachtest/tests/decommissionbench.go
@@ -272,7 +272,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")
 	}
diff --git a/pkg/cmd/roachtest/tests/follower_reads.go b/pkg/cmd/roachtest/tests/follower_reads.go
index 2290c6cff98b..7a9b70854712 100644
--- a/pkg/cmd/roachtest/tests/follower_reads.go
+++ b/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) {
diff --git a/pkg/cmd/roachtest/tests/import.go b/pkg/cmd/roachtest/tests/import.go
index 1caa658969b9..795af85c08fe 100644
--- a/pkg/cmd/roachtest/tests/import.go
+++ b/pkg/cmd/roachtest/tests/import.go
@@ -189,8 +189,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,
diff --git a/pkg/cmd/roachtest/tests/indexes.go b/pkg/cmd/roachtest/tests/indexes.go
index 58b0eb3e21c8..83d6cf8af93b 100644
--- a/pkg/cmd/roachtest/tests/indexes.go
+++ b/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)
diff --git a/pkg/cmd/roachtest/tests/ledger.go b/pkg/cmd/roachtest/tests/ledger.go
index cbd164ecf39e..a78387ca7b66 100644
--- a/pkg/cmd/roachtest/tests/ledger.go
+++ b/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)
diff --git a/pkg/cmd/roachtest/tests/mixed_version_cdc.go b/pkg/cmd/roachtest/tests/mixed_version_cdc.go
index 8285567dcffc..1f0a8cddf37e 100644
--- a/pkg/cmd/roachtest/tests/mixed_version_cdc.go
+++ b/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:          30 * 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) {
diff --git a/pkg/cmd/roachtest/tests/restore.go b/pkg/cmd/roachtest/tests/restore.go
index 8cff79545767..a9502faec26d 100644
--- a/pkg/cmd/roachtest/tests/restore.go
+++ b/pkg/cmd/roachtest/tests/restore.go
@@ -481,7 +481,11 @@ func (hw hardwareSpecs) makeClusterSpecs(r registry.Registry, backupCloud string
 		clusterOpts = append(clusterOpts, spec.Mem(hw.mem))
 	}
 	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 {
diff --git a/pkg/cmd/roachtest/tests/roachmart.go b/pkg/cmd/roachtest/tests/roachmart.go
index bfbe3c7925eb..1a7c6b891d5f 100644
--- a/pkg/cmd/roachtest/tests/roachmart.go
+++ b/pkg/cmd/roachtest/tests/roachmart.go
@@ -73,8 +73,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) {
diff --git a/pkg/cmd/roachtest/tests/schemachange_random_load.go b/pkg/cmd/roachtest/tests/schemachange_random_load.go
index 6a821e6c4af0..da116cc4bef2 100644
--- a/pkg/cmd/roachtest/tests/schemachange_random_load.go
+++ b/pkg/cmd/roachtest/tests/schemachange_random_load.go
@@ -39,11 +39,6 @@ type randomLoadBenchSpec struct {
 }
 
 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,
@@ -51,8 +46,10 @@ func registerSchemaChangeRandomLoad(r registry.Registry) {
 		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,
diff --git a/pkg/cmd/roachtest/tests/tpcc.go b/pkg/cmd/roachtest/tests/tpcc.go
index b68dcd7f52f3..8104856bf9b8 100644
--- a/pkg/cmd/roachtest/tests/tpcc.go
+++ b/pkg/cmd/roachtest/tests/tpcc.go
@@ -636,8 +636,8 @@ func registerTPCC(r registry.Registry) {
 				Name:  tc.name,
 				Owner: registry.OwnerSQLFoundations,
 				// Add an extra node which serves as the workload nodes.
-				Cluster:           r.MakeClusterSpec(len(regions)*nodesPerRegion+1, spec.Geo(), spec.Zones(strings.Join(zs, ","))),
-				CompatibleClouds:  registry.AllExceptAWS,
+				Cluster:           r.MakeClusterSpec(len(regions)*nodesPerRegion+1, spec.Geo(), spec.GCEZones(strings.Join(zs, ","))),
+				CompatibleClouds:  registry.OnlyGCE,
 				Suites:            registry.Suites(registry.Nightly),
 				EncryptionSupport: registry.EncryptionMetamorphic,
 				Leases:            registry.MetamorphicLeases,
@@ -830,7 +830,7 @@ func registerTPCC(r registry.Registry) {
 		EstimatedMaxGCE:   5000,
 		EstimatedMaxAWS:   5000,
 
-		Clouds: registry.AllExceptAWS,
+		Clouds: registry.OnlyGCE,
 		Suites: registry.Suites(registry.Nightly),
 	})
 	registerTPCCBenchSpec(r, tpccBenchSpec{
@@ -845,7 +845,7 @@ func registerTPCC(r registry.Registry) {
 		EstimatedMaxGCE:   2000,
 		EstimatedMaxAWS:   2000,
 
-		Clouds: registry.AllExceptAWS,
+		Clouds: registry.OnlyGCE,
 		Suites: registry.Suites(registry.Nightly),
 	})
 	registerTPCCBenchSpec(r, tpccBenchSpec{
@@ -1110,10 +1110,10 @@ func registerTPCCBenchSpec(r registry.Registry, b tpccBenchSpec) {
 		// No specifier.
 	case multiZone:
 		nameParts = append(nameParts, "multi-az")
-		opts = append(opts, spec.Geo(), spec.Zones(strings.Join(b.Distribution.zones(), ",")))
+		opts = append(opts, spec.Geo(), spec.GCEZones(strings.Join(b.Distribution.zones(), ",")))
 	case multiRegion:
 		nameParts = append(nameParts, "multi-region")
-		opts = append(opts, spec.Geo(), spec.Zones(strings.Join(b.Distribution.zones(), ",")))
+		opts = append(opts, spec.Geo(), spec.GCEZones(strings.Join(b.Distribution.zones(), ",")))
 	default:
 		panic("unexpected")
 	}

From 31718a3a99b8f80c69a5f238a2558aadb6449f8c Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Wed, 4 Oct 2023 13:38:35 -0700
Subject: [PATCH 5/5] roachtest: stop using ClusterSpec.Cloud in test code

This change removes all remaining uses of `ClusterSpec.Cloud` except
those internal to roachtest. Code that is part of running a test now
uses `Cluster.Cloud()` instead.

Informs: #104029
Release note: None
---
 pkg/cmd/roachtest/tests/awsdms.go             |  2 +-
 pkg/cmd/roachtest/tests/backup.go             |  6 +++---
 pkg/cmd/roachtest/tests/copy.go               |  2 +-
 pkg/cmd/roachtest/tests/disk_stall.go         | 14 ++++++-------
 pkg/cmd/roachtest/tests/follower_reads.go     |  2 +-
 pkg/cmd/roachtest/tests/import.go             |  6 +++---
 .../roachtest/tests/import_cancellation.go    |  2 +-
 pkg/cmd/roachtest/tests/indexes.go            |  2 +-
 pkg/cmd/roachtest/tests/kv.go                 | 21 ++++++++++---------
 .../roachtest/tests/mixed_version_backup.go   |  2 +-
 .../mixed_version_decl_schemachange_compat.go |  2 +-
 pkg/cmd/roachtest/tests/rebalance_load.go     |  7 +------
 pkg/cmd/roachtest/tests/restore.go            |  3 +--
 pkg/cmd/roachtest/tests/schemachange.go       |  2 +-
 pkg/cmd/roachtest/tests/sqlsmith.go           |  2 +-
 pkg/cmd/roachtest/tests/tpc_utils.go          |  2 +-
 pkg/cmd/roachtest/tests/tpcdsvec.go           |  2 +-
 pkg/cmd/roachtest/tests/ycsb.go               |  2 +-
 18 files changed, 38 insertions(+), 43 deletions(-)

diff --git a/pkg/cmd/roachtest/tests/awsdms.go b/pkg/cmd/roachtest/tests/awsdms.go
index d15a2a4d38d2..aff716d80f0d 100644
--- a/pkg/cmd/roachtest/tests/awsdms.go
+++ b/pkg/cmd/roachtest/tests/awsdms.go
@@ -210,7 +210,7 @@ func runAWSDMS(ctx context.Context, t test.Test, c cluster.Cluster) {
 		t.Fatal("cannot be run in local mode")
 	}
 	// We may not have the requisite certificates to start DMS/RDS on non-AWS invocations.
-	if cloud := c.Spec().Cloud; cloud != spec.AWS {
+	if cloud := c.Cloud(); cloud != spec.AWS {
 		t.Skipf("skipping test on cloud %s", cloud)
 		return
 	}
diff --git a/pkg/cmd/roachtest/tests/backup.go b/pkg/cmd/roachtest/tests/backup.go
index 9cd279d24b50..ea1a4cb3181e 100644
--- a/pkg/cmd/roachtest/tests/backup.go
+++ b/pkg/cmd/roachtest/tests/backup.go
@@ -444,7 +444,7 @@ func registerBackup(r registry.Registry) {
 			CompatibleClouds:  registry.AllExceptAWS,
 			Suites:            registry.Suites(registry.Nightly),
 			Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-				if c.Spec().Cloud != item.machine {
+				if c.Cloud() != item.machine {
 					t.Skip("backup assumeRole is only configured to run on "+item.machine, "")
 				}
 
@@ -555,7 +555,7 @@ func registerBackup(r registry.Registry) {
 			Suites:            registry.Suites(registry.Nightly),
 			Tags:              item.tags,
 			Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-				if c.Spec().Cloud != item.machine {
+				if c.Cloud() != item.machine {
 					t.Skip("backupKMS roachtest is only configured to run on "+item.machine, "")
 				}
 
@@ -899,7 +899,7 @@ func registerBackup(r registry.Registry) {
 		CompatibleClouds:  registry.AllExceptAWS,
 		Suites:            registry.Suites(registry.Nightly),
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-			if c.Spec().Cloud != spec.GCE {
+			if c.Cloud() != spec.GCE {
 				t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 			}
 			runBackupMVCCRangeTombstones(ctx, t, c, mvccRangeTombstoneConfig{})
diff --git a/pkg/cmd/roachtest/tests/copy.go b/pkg/cmd/roachtest/tests/copy.go
index c3d83b3e47f7..f56e449cc127 100644
--- a/pkg/cmd/roachtest/tests/copy.go
+++ b/pkg/cmd/roachtest/tests/copy.go
@@ -199,7 +199,7 @@ func registerCopy(r registry.Registry) {
 			Suites:           registry.Suites(registry.Nightly),
 			Leases:           registry.MetamorphicLeases,
 			Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-				if c.Spec().Cloud != spec.GCE {
+				if c.Cloud() != spec.GCE {
 					t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 				}
 				runCopy(ctx, t, c, tc.rows, tc.txn)
diff --git a/pkg/cmd/roachtest/tests/disk_stall.go b/pkg/cmd/roachtest/tests/disk_stall.go
index 5c70a99ec96e..643a206fdeca 100644
--- a/pkg/cmd/roachtest/tests/disk_stall.go
+++ b/pkg/cmd/roachtest/tests/disk_stall.go
@@ -290,7 +290,7 @@ type dmsetupDiskStaller struct {
 
 var _ diskStaller = (*dmsetupDiskStaller)(nil)
 
-func (s *dmsetupDiskStaller) device() string { return getDevice(s.t, s.c.Spec()) }
+func (s *dmsetupDiskStaller) device() string { return getDevice(s.t, s.c) }
 
 func (s *dmsetupDiskStaller) Setup(ctx context.Context) {
 	dev := s.device()
@@ -365,15 +365,15 @@ func (s *cgroupDiskStaller) Unstall(ctx context.Context, nodes option.NodeListOp
 func (s *cgroupDiskStaller) device() (major, minor int) {
 	// TODO(jackson): Programmatically determine the device major,minor numbers.
 	// eg,:
-	//    deviceName := getDevice(s.t, s.c.Spec())
+	//    deviceName := getDevice(s.t, s.c)
 	//    `cat /proc/partitions` and find `deviceName`
-	switch s.c.Spec().Cloud {
+	switch s.c.Cloud() {
 	case spec.GCE:
 		// ls -l /dev/sdb
 		// brw-rw---- 1 root disk 8, 16 Mar 27 22:08 /dev/sdb
 		return 8, 16
 	default:
-		s.t.Fatalf("unsupported cloud %q", s.c.Spec().Cloud)
+		s.t.Fatalf("unsupported cloud %q", s.c.Cloud())
 		return 0, 0
 	}
 }
@@ -391,14 +391,14 @@ func (s *cgroupDiskStaller) setThroughput(
 	))
 }
 
-func getDevice(t test.Test, s spec.ClusterSpec) string {
-	switch s.Cloud {
+func getDevice(t test.Test, c cluster.Cluster) string {
+	switch c.Cloud() {
 	case spec.GCE:
 		return "/dev/sdb"
 	case spec.AWS:
 		return "/dev/nvme1n1"
 	default:
-		t.Fatalf("unsupported cloud %q", s.Cloud)
+		t.Fatalf("unsupported cloud %q", c.Cloud())
 		return ""
 	}
 }
diff --git a/pkg/cmd/roachtest/tests/follower_reads.go b/pkg/cmd/roachtest/tests/follower_reads.go
index 7a9b70854712..eb06ee596519 100644
--- a/pkg/cmd/roachtest/tests/follower_reads.go
+++ b/pkg/cmd/roachtest/tests/follower_reads.go
@@ -68,7 +68,7 @@ func registerFollowerReads(r registry.Registry) {
 			Suites:           registry.Suites(registry.Nightly),
 			Leases:           registry.MetamorphicLeases,
 			Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-				if c.Spec().Cloud == spec.GCE && c.Spec().Arch == vm.ArchARM64 {
+				if c.Cloud() == spec.GCE && c.Spec().Arch == vm.ArchARM64 {
 					t.Skip("arm64 in GCE is available only in us-central1")
 				}
 				c.Start(ctx, t.L(), option.DefaultStartOpts(), install.MakeClusterSettings())
diff --git a/pkg/cmd/roachtest/tests/import.go b/pkg/cmd/roachtest/tests/import.go
index 795af85c08fe..8d5ad0113f67 100644
--- a/pkg/cmd/roachtest/tests/import.go
+++ b/pkg/cmd/roachtest/tests/import.go
@@ -92,7 +92,7 @@ func registerImportNodeShutdown(r registry.Registry) {
 		Suites:           registry.Suites(registry.Nightly),
 		Leases:           registry.MetamorphicLeases,
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-			if c.Spec().Cloud != spec.GCE {
+			if c.Cloud() != spec.GCE {
 				t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 			}
 			c.Start(ctx, t.L(), option.DefaultStartOpts(), install.MakeClusterSettings())
@@ -111,7 +111,7 @@ func registerImportNodeShutdown(r registry.Registry) {
 		Suites:           registry.Suites(registry.Nightly),
 		Leases:           registry.MetamorphicLeases,
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-			if c.Spec().Cloud != spec.GCE {
+			if c.Cloud() != spec.GCE {
 				t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 			}
 			c.Start(ctx, t.L(), option.DefaultStartOpts(), install.MakeClusterSettings())
@@ -232,7 +232,7 @@ func registerImportTPCH(r registry.Registry) {
 			EncryptionSupport: registry.EncryptionMetamorphic,
 			Leases:            registry.MetamorphicLeases,
 			Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-				if c.Spec().Cloud != spec.GCE {
+				if c.Cloud() != spec.GCE {
 					t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 				}
 				tick, perfBuf := initBulkJobPerfArtifacts(t.Name(), item.timeout)
diff --git a/pkg/cmd/roachtest/tests/import_cancellation.go b/pkg/cmd/roachtest/tests/import_cancellation.go
index e3613c9d7726..bddc29ae0ddb 100644
--- a/pkg/cmd/roachtest/tests/import_cancellation.go
+++ b/pkg/cmd/roachtest/tests/import_cancellation.go
@@ -40,7 +40,7 @@ func registerImportCancellation(r registry.Registry) {
 			Suites:           registry.Suites(registry.Nightly),
 			Leases:           registry.MetamorphicLeases,
 			Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-				if c.Spec().Cloud != spec.GCE {
+				if c.Cloud() != spec.GCE {
 					t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 				}
 				runImportCancellation(ctx, t, c, rangeTombstones)
diff --git a/pkg/cmd/roachtest/tests/indexes.go b/pkg/cmd/roachtest/tests/indexes.go
index 83d6cf8af93b..5bdaebed30d5 100644
--- a/pkg/cmd/roachtest/tests/indexes.go
+++ b/pkg/cmd/roachtest/tests/indexes.go
@@ -47,7 +47,7 @@ func registerNIndexes(r registry.Registry, secondaryIndexes int) {
 		// Uses CONFIGURE ZONE USING ... COPY FROM PARENT syntax.
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
 			firstAZ := gceGeoZones[0]
-			if c.Spec().Cloud == spec.AWS {
+			if c.Cloud() == spec.AWS {
 				firstAZ = awsGeoZones[0]
 			}
 			roachNodes := c.Range(1, nodes)
diff --git a/pkg/cmd/roachtest/tests/kv.go b/pkg/cmd/roachtest/tests/kv.go
index d52f900a0322..b134477aa4d0 100644
--- a/pkg/cmd/roachtest/tests/kv.go
+++ b/pkg/cmd/roachtest/tests/kv.go
@@ -307,25 +307,26 @@ func registerKV(r registry.Registry) {
 		}
 		cSpec := r.MakeClusterSpec(opts.nodes+1, spec.CPU(opts.cpus), spec.SSD(opts.ssds), spec.RAID0(opts.raid0))
 
-		clouds := registry.AllExceptAWS
-		var tags map[string]struct{}
-		// All the kv0|95 tests should run on AWS by default
-		if !opts.weekly && opts.ssds == 0 && (opts.readPercent == 95 || opts.readPercent == 0) {
+		var clouds registry.CloudSet
+		tags := make(map[string]struct{})
+		if opts.ssds != 0 {
+			// Multi-store tests are only supported on GCE.
+			clouds = registry.OnlyGCE
+		} else if !opts.weekly && (opts.readPercent == 95 || opts.readPercent == 0) {
+			// All the kv0|95 tests should run on AWS.
 			clouds = registry.AllClouds
 			tags = registry.Tags("aws")
+		} else {
+			clouds = registry.AllExceptAWS
 		}
+
 		suites := registry.Suites(registry.Nightly)
 		if opts.weekly {
 			suites = registry.Suites(registry.Weekly)
-			tags = registry.Tags("weekly")
+			tags["weekly"] = struct{}{}
 		}
 
-		var skip string
-		if opts.ssds != 0 && cSpec.Cloud != spec.GCE {
-			skip = fmt.Sprintf("multi-store tests are not supported on cloud %s", cSpec.Cloud)
-		}
 		r.Add(registry.TestSpec{
-			Skip:      skip,
 			Name:      strings.Join(nameParts, "/"),
 			Owner:     owner,
 			Benchmark: true,
diff --git a/pkg/cmd/roachtest/tests/mixed_version_backup.go b/pkg/cmd/roachtest/tests/mixed_version_backup.go
index a2593a6e85a2..a5c4fc89ba2e 100644
--- a/pkg/cmd/roachtest/tests/mixed_version_backup.go
+++ b/pkg/cmd/roachtest/tests/mixed_version_backup.go
@@ -2468,7 +2468,7 @@ func registerBackupMixedVersion(r registry.Registry) {
 		CompatibleClouds:  registry.AllExceptAWS,
 		Suites:            registry.Suites(registry.Nightly),
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-			if c.Spec().Cloud != spec.GCE && !c.IsLocal() {
+			if c.Cloud() != spec.GCE && !c.IsLocal() {
 				t.Skip("uses gs://cockroachdb-backup-testing-long-ttl; see https://github.com/cockroachdb/cockroach/issues/105968")
 			}
 
diff --git a/pkg/cmd/roachtest/tests/mixed_version_decl_schemachange_compat.go b/pkg/cmd/roachtest/tests/mixed_version_decl_schemachange_compat.go
index 8d3404f2523d..b650c70f8a79 100644
--- a/pkg/cmd/roachtest/tests/mixed_version_decl_schemachange_compat.go
+++ b/pkg/cmd/roachtest/tests/mixed_version_decl_schemachange_compat.go
@@ -34,7 +34,7 @@ func registerDeclSchemaChangeCompatMixedVersions(r registry.Registry) {
 		CompatibleClouds: registry.AllExceptAWS,
 		Suites:           registry.Suites(registry.Nightly),
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-			if c.Spec().Cloud != spec.GCE {
+			if c.Cloud() != spec.GCE {
 				t.Skip("uses gs://cockroach-corpus; see https://github.com/cockroachdb/cockroach/issues/105968")
 			}
 			runDeclSchemaChangeCompatMixedVersions(ctx, t, c)
diff --git a/pkg/cmd/roachtest/tests/rebalance_load.go b/pkg/cmd/roachtest/tests/rebalance_load.go
index 25079acc03be..3b6506b4c75c 100644
--- a/pkg/cmd/roachtest/tests/rebalance_load.go
+++ b/pkg/cmd/roachtest/tests/rebalance_load.go
@@ -283,17 +283,12 @@ func registerRebalanceLoad(r registry.Registry) {
 		},
 	)
 	cSpec := r.MakeClusterSpec(7, spec.SSD(2)) // the last node is just used to generate load
-	var skip string
-	if cSpec.Cloud != spec.GCE {
-		skip = fmt.Sprintf("multi-store tests are not supported on cloud %s", cSpec.Cloud)
-	}
 	r.Add(
 		registry.TestSpec{
-			Skip:             skip,
 			Name:             `rebalance/by-load/replicas/ssds=2`,
 			Owner:            registry.OwnerKV,
 			Cluster:          cSpec,
-			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) {
diff --git a/pkg/cmd/roachtest/tests/restore.go b/pkg/cmd/roachtest/tests/restore.go
index a9502faec26d..216b9b1ee57f 100644
--- a/pkg/cmd/roachtest/tests/restore.go
+++ b/pkg/cmd/roachtest/tests/restore.go
@@ -745,8 +745,7 @@ func makeRestoreDriver(t test.Test, c cluster.Cluster, sp restoreSpecs) restoreD
 }
 
 func (rd *restoreDriver) prepareCluster(ctx context.Context) {
-
-	if rd.c.Spec().Cloud != rd.sp.backup.cloud {
+	if rd.c.Cloud() != rd.sp.backup.cloud {
 		// For now, only run the test on the cloud provider that also stores the backup.
 		rd.t.Skipf("test configured to run on %s", rd.sp.backup.cloud)
 	}
diff --git a/pkg/cmd/roachtest/tests/schemachange.go b/pkg/cmd/roachtest/tests/schemachange.go
index a2236658d0b1..4df4ce9cd8e5 100644
--- a/pkg/cmd/roachtest/tests/schemachange.go
+++ b/pkg/cmd/roachtest/tests/schemachange.go
@@ -36,7 +36,7 @@ func registerSchemaChangeDuringKV(r registry.Registry) {
 		Suites:           registry.Suites(registry.Nightly),
 		Leases:           registry.MetamorphicLeases,
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-			if c.Spec().Cloud != spec.GCE {
+			if c.Cloud() != spec.GCE {
 				t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 			}
 			const fixturePath = `gs://cockroach-fixtures/workload/tpch/scalefactor=10/backup?AUTH=implicit`
diff --git a/pkg/cmd/roachtest/tests/sqlsmith.go b/pkg/cmd/roachtest/tests/sqlsmith.go
index be8a0c820d64..f7742a7447c1 100644
--- a/pkg/cmd/roachtest/tests/sqlsmith.go
+++ b/pkg/cmd/roachtest/tests/sqlsmith.go
@@ -296,7 +296,7 @@ WITH into_db = 'defaultdb', unsafe_restore_incompatible_version;
 			// NB: sqlsmith failures should never block a release.
 			NonReleaseBlocker: true,
 			Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-				if c.Spec().Cloud != spec.GCE {
+				if c.Cloud() != spec.GCE {
 					t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 				}
 				runSQLSmith(ctx, t, c, setup, setting)
diff --git a/pkg/cmd/roachtest/tests/tpc_utils.go b/pkg/cmd/roachtest/tests/tpc_utils.go
index ebf613d10f48..78887dd20550 100644
--- a/pkg/cmd/roachtest/tests/tpc_utils.go
+++ b/pkg/cmd/roachtest/tests/tpc_utils.go
@@ -44,7 +44,7 @@ func loadTPCHDataset(
 	disableMergeQueue bool,
 	secure bool,
 ) (retErr error) {
-	if c.Spec().Cloud != spec.GCE {
+	if c.Cloud() != spec.GCE {
 		t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 	}
 
diff --git a/pkg/cmd/roachtest/tests/tpcdsvec.go b/pkg/cmd/roachtest/tests/tpcdsvec.go
index 54e6573c0a05..31e5000a2b9b 100644
--- a/pkg/cmd/roachtest/tests/tpcdsvec.go
+++ b/pkg/cmd/roachtest/tests/tpcdsvec.go
@@ -192,7 +192,7 @@ WITH unsafe_restore_incompatible_version;
 		CompatibleClouds: registry.AllExceptAWS,
 		Suites:           registry.Suites(registry.Nightly),
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-			if c.Spec().Cloud != spec.GCE {
+			if c.Cloud() != spec.GCE {
 				t.Skip("uses gs://cockroach-fixtures; see https://github.com/cockroachdb/cockroach/issues/105968")
 			}
 			runTPCDSVec(ctx, t, c)
diff --git a/pkg/cmd/roachtest/tests/ycsb.go b/pkg/cmd/roachtest/tests/ycsb.go
index ede2e9ca21c7..538c37df534e 100644
--- a/pkg/cmd/roachtest/tests/ycsb.go
+++ b/pkg/cmd/roachtest/tests/ycsb.go
@@ -49,7 +49,7 @@ func registerYCSB(r registry.Registry) {
 	) {
 		// For now, we only want to run the zfs tests on GCE, since only GCE supports
 		// starting roachprod instances on zfs.
-		if c.Spec().FileSystem == spec.Zfs && c.Spec().Cloud != spec.GCE {
+		if c.Spec().FileSystem == spec.Zfs && c.Cloud() != spec.GCE {
 			t.Skip("YCSB zfs benchmark can only be run on GCE", "")
 		}