Merge #153269 #153964 #154476

153269: roachtest/mixedversion: sync workload binary versions with the cluster version r=DarrylWong a=williamchoe3

Fixes #147374 

### Problem
The current workload binary is no longer backwards compatible with certain workloads i.e. `bank` (see original issue above for more information). Previously, the workload binary would just use whatever default binary was staged on the workload node by default. In a test, the test author could explicitly download an older version of workload to circumvent version compatibility issues or use `Workload`'s `override` ([PR](https://github.com/cockroachdb/cockroach/pull/151165/files#diff-11eb39229832689838a549c5cd62b0adeb799fc4e58468ab17268ffcc86c6bd4L829 )) parameter to download a binary matching the cluster version to the workload node when the workload hook is being executed.
#### Change
While the later's solution works, it makes more sense to always just use a workload binary that matches the cluster. If a test wants to use a specific workload binary version though, that functionality remains possible as well.
During `startStep`, while we are staging the cluster at it's initial version, we also stage ~~that version's binary~~ all the versioned binaries in the upgrade plan on the workload node(s). Previously we just left the workload node alone at this step.  ~~Then, introduced a new test framework specific step `stageWorkloadBinaryStep` which gets added in `test.Planner.afterUpgradeSteps` which get added to the plan after a cluster finalizes it's upgrade.  `stageWorkloadBinaryStep` stages the finalized cluster version's binary on the workload node(s).~~ 

Removed the call to `UploadCockroach` from `Workload` as when the workload init or run hooks are called, the cluster binary will already exist on the workload node. ~~Kept `override` parameter, but now if override is set, `Workload` will simply just respect the binary that get's passed into the command instead of defaulting the current cluster version.~~
* If a test wants to use a specific Workload binary, they are able to do that if they want without using the Framework's helper function

Added a new helper method to `Helper` which can be used in hooks to abstract a lower level roachtest framework call
```go
clusterupgrade.BinaryPathForVersion(t.rt, h.System.FromVersion, "cockroach")
--> 
CockroachBinaryForWorkload(t)
```

Added new mixedversion test option which will add the framework step to stage all versioned cockroach binaries in the upgrade plan to the WorkloadNode
```go
mixedversion.WithWorkloadNodes(c.WorkloadNode())
```

### Example
Given a plan e.g. during the `run startup hooks concurrently` set of steps
```
Plan:
Upgrades:           v25.2.5 → <current>
Deployment mode:    system-only
Mutators:           cluster_setting[kv.transaction.write_buffering.enabled], cluster_setting[kv.rangefeed.buffered_sender.enabled]
Plan:
├── install fixtures for version "v25.2.5" (1)
├── start cluster at version "v25.2.5" (2)
├── wait for all nodes (:1-4) to acknowledge cluster version '25.2' on system tenant (3)
├── stage workload binary on workload node(s) :5 for version(s) v25.2.5, <current> (4) <-- NEW
├── run startup hooks concurrently
│   ├── run "maybe enable tenant features", after 30s delay (5)
│   ├── run "load TPCC dataset", after 500ms delay (6)
│   ├── run "load bank dataset", after 30s delay (7)
│   └── set cluster setting "kv.rangefeed.buffered_sender.enabled" to 'true' on system tenant, after 3m0s delay (8)
└── upgrade cluster from "v25.2.5" to "<current>"
```
### Verification
Adhoc Nightly against all mixed version tests https://teamcity.cockroachdb.com/buildConfiguration/Cockroach_Nightlies_RoachtestNightlyGceBazel/20482572?buildTab=overview&expandBuildDeploymentsSection=false&expandBuildChangesSection=true&hideTestsFromDependencies=false&expandBuildTestsSection=true&hideProblemsFromDependencies=false&expandBuildProblemsSection=true

Added a `gomock` for Cluster in cluster/mock for unit tests
* Removed duplicate Cluster mock in clusterstats
* Wanted to colocate the mocks with the actual interface they are mocking. Then each package that wants to use the mock just needs to import that mock instead of creating a new mock in that package and creating a new generated file for that mock
* 
Added datadriven unit test that captures the new post upgrade internal framework step

While looking at `Test` I saw these test specific fields, (not sure what this specific methodology is called), but just curious  because my initial thought was to add a mock Cluster implementation, because I have workload node related logic that needs to call cluster interface methods. I guess i could've maybe went with something like below, but I personally don't like intertwining prod logic and test logic. I guess is there any reason we wouldn't want the mock cluster? Besides the boiler plate it creates in the unit tests?
 https://github.com/cockroachdb/cockroach/blob/b4738823bcc55547ff431b99684576370760e56d/pkg/cmd/roachtest/roachtestutil/mixedversion/planner_test.go

```
	// Test is the main struct callers of this package interact with.
	Test struct {
		// the following are test-only fields, allowing tests to simulate
		// cluster properties without passing a cluster.Cluster
		// implementation.
		_arch      *vm.CPUArch
		_isLocal   *bool
		_getFailer func(name string) (*failures.Failer, error)
```  

### mixedversion test refactor
Key
👌: no workload, no change needed
➕: added mixedversion.WithWorkloadNodes(c.WorkloadNode()) test option to stage versioned binaries on workload node 
✅: removed hardcoded "./cockroach workload ..." call to something like "%s workload... ", h.CockroachBinaryForWorkload(t)
⚠️: can't add versioned call because workload logic not defined in a user hook. 
  * Most of these tests use Workload() which will override the binary with the versioned one by default
  * Some of these tests have logic around the Run cmd so calling Workload doesn't provide an easy fix, but the test should continue to work with the unversioned (default) binary

❌: not refactoring bc it's not using bank and the logic is tighty coupled with shared non mixedversion roachtests so would need to decouple, which I think would make the test logic harder to understand and not worth having a 1 line helper 

Test Changes
```
👌acceptance/validate-system-schema-after-version-upgrade [sql-foundations]   randomized,timeout: 1h0m0s
👌acceptance/version-upgrade [test-eng]   randomized,timeout: 2h0m0s
➕✅admission-control/elastic-workload/mixed-version [kv]   timeout: 3h0m0s
➕⚠️backup-restore/mixed-version [disaster-recovery]   randomized,timeout: 8h0m0s
➕⚠️c2c/mixed-version [disaster-recovery]
  * 2 clusters, 1 workload node, put the test option to upgrade the workload node in both clusters, techniaclly 1 of them might be redundant, but i just wanna be sure by the time either cluster makes a workload call it's good, not sure if it's the 1st cluster, 2nd cluster, or random which one goes first so just putting in both, also binary existence is checked before staging so no redundent downloads will be done
➕✅cdc/mixed-version/checkpointing [cdc]   randomized,timeout: 3h0m0s
➕✅cdc/mixed-versions [cdc]   randomized,timeout: 3h0m0s
👌change-replicas/mixed-version [kv]   randomized,timeout: 1h0m0s
👌db-console/mixed-version-cypress [obs-prs]   timeout: 2h0m0s
* uses workload node for another task that's not running workload so no change is needed
➕✅db-console/mixed-version-endpoints [obs-prs]   randomized,timeout: 1h0m0s
👌declarative_schema_changer/job-compatibility-mixed-version-V242-V243 [sql-foundations]   randomized
👌~~➕✅~~decommission/mixed-versions [kv]   randomized
  * ~~this test didn't have a workload node previously and I didn't see a reason not to use one (unlike some other tests that want the extra pressure on the node)~~
  * data loading is fine to do on a cluster node (see comments below)
👌follower-reads/mixed-version/single-region [kv]   randomized
👌follower-reads/mixed-version/survival=region/locality=global/reads=strong [kv]   randomized
👌http-register-routes/mixed-version [obs-prs]   randomized,timeout: 1h0m0s
~~➕✅~~import/mixed-versions [sql-queries] (skipped: Issue #143870)   randomized
* don't need to use an explicit workload node here (see comments below)
👌jobs/mixed-versions [disaster-recovery]   randomized
➕✅⚠️ldr/mixed-version [disaster-recovery]
  * the init is in a hook, i can user the helper, the run isn't in a hook it uses Workload, so i didn't replace it
➕❌multi-region/mixed-version [test-eng]   randomized,timeout: 36h0m0s
❌multitenant-upgrade [server]   randomized,timeout: 5h0m0s
❌rebalance/by-load/leases/mixed-version [kv]   randomized
❌rebalance/by-load/replicas/mixed-version [kv]   randomized
❌schemachange/mixed-versions [sql-foundations]   randomized
❌schemachange/mixed-versions-compat [sql-foundations]
❌schemachange/secondary-index-multi-version [sql-foundations]   randomized
➕⚠️sql-stats/mixed-version [obs-prs]   randomized,timeout: 1h0m0s
➕✅tpcc/mixed-headroom/chaos/n6cpu16 [test-eng]   randomized,timeout: 7h0m0s
➕✅tpcc/mixed-headroom/n5cpu16 [test-eng]   randomized,timeout: 7h0m0s
👌validate-system-schema-after-version-upgrade/separate-process [sql-foundations]
```


#### Notes
`pkg/cmd/roachtest/clusterstats/BUILD.bazel` I made some edits initially, removed them, and now ./dev gen is generating go_library / go_test deps in a different order 🤷‍♂️

* Saw in `pkg/cmd/roachtest/roachtestutil/clusterupgrade/clusterupgrade.go` in `clusterupgrade` pkg, `UploadWorkload` but chose not to use it because it's a currently unused path that also eventually calls a deprecated Cluster API function and `UploadCockroach` & `UploadWorkload` are both just wrappers around `uploadBinaryVersion`



153964: roachtest: allow specifying a set of valid architectures r=herkolategan a=DarrylWong

Previously, tests could only specify one valid architecture or enable all of them to be metamorphically chosen. This changes the cluster spec to allow multiple architectures to be specified.

This is now needed because we recently started running FIPS metamorphically on our nightlies. Some tests are not compatible with ARM or FIPS, and we want to only disable execution on those architectures.

Fixes: #153833
Fixes: #154110
Informs: #152781

154476: build: update patched Go to the latest commit r=rail,RaduBerinde a=rickystewart

To consume cockroachdb/go@889a976.

Co-authored-by: William Choe <[email protected]>
Co-authored-by: DarrylWong <[email protected]>
Co-authored-by: Ricky Stewart <[email protected]>

Author: 4 people

WORKSPACE

@@ -167,14 +167,14 @@ load(
 go_download_sdk(
     name = "go_sdk",
     sdks = {
-        "darwin_amd64": ("go1.23.12.darwin-amd64.tar.gz", "34457131f14281e21e25493d68e7519ccf26342d176dac36a4fc5dbf5ef122d9"),
-        "darwin_arm64": ("go1.23.12.darwin-arm64.tar.gz", "30e0735ab9ccda203946536d24afe895abd1a1d3f35ad199f9768ccbdd5d60bc"),
-        "linux_amd64": ("go1.23.12.linux-amd64.tar.gz", "0cac0ac930ecb9458b8a0a7969cbf735c5884d24c879c97eb28a8997eca986fa"),
-        "linux_arm64": ("go1.23.12.linux-arm64.tar.gz", "528601fc8fb2c7e5ce8b7ae7651fd4fce2450bbef687beb96616edc5a9effa41"),
-        "linux_s390x": ("go1.23.12.linux-s390x.tar.gz", "f3f11bbb731da6716776d1c29a2db3d1063fa0a9f8c00636e6a77793ba79e2e3"),
-        "windows_amd64": ("go1.23.12.windows-amd64.tar.gz", "71b5b5b86b3a5ff9f124e21984abd874a6bfeb438f368de2eee7c60a25a19c94"),
+        "darwin_amd64": ("go1.23.12.darwin-amd64.tar.gz", "25b853b77448c6f196475c1ab44d4a617b719db34e9e072c72c927a905bdbce6"),
+        "darwin_arm64": ("go1.23.12.darwin-arm64.tar.gz", "fae1c6d45f72d559b9fe52224918321e7642cb75cd6a71186c7ccdf577189a4f"),
+        "linux_amd64": ("go1.23.12.linux-amd64.tar.gz", "03ac2c00dfde86a4beb3e94f063edffe1d2c089dfbbe1ec2776d1531c58e196b"),
+        "linux_arm64": ("go1.23.12.linux-arm64.tar.gz", "442f8d3510d141d6f479f05f7de00560aa1df59e0a5e1d02ca58ed9e4757d0d7"),
+        "linux_s390x": ("go1.23.12.linux-s390x.tar.gz", "6e3c09786b434bdd8d44de714656bef7c4fe484d2a2e6f8c98a797385e5f280e"),
+        "windows_amd64": ("go1.23.12.windows-amd64.tar.gz", "513af02afaa1c64501f1ec0a9c9445dbf88ea55c0b7a5a86e0b7fc94f0c59b39"),
     },
-    urls = ["https://storage.googleapis.com/public-bazel-artifacts/go/20250818-202337/{}"],
+    urls = ["https://storage.googleapis.com/public-bazel-artifacts/go/20250930-204932/{}"],
     version = "1.23.12",
 )
 
@@ -659,9 +659,8 @@ go_download_sdk(
     # able to provide additional diagnostic information such as the expected version of OpenSSL.
     experiments = ["boringcrypto"],
     sdks = {
-
-        "linux_amd64": ("go1.23.12fips.linux-amd64.tar.gz", "9c58fd7137b4c9d387a5c37fd2e728bc5d39357c7f8ba3358bcae513704c2983"),
+        "linux_amd64": ("go1.23.12fips.linux-amd64.tar.gz", "ae2d57fa43ef68aa70e6c0c1def065bd2806411acfc0d59591ddff81168be095"),
     },
-    urls = ["https://storage.googleapis.com/public-bazel-artifacts/go/20250818-202337/{}"],
+    urls = ["https://storage.googleapis.com/public-bazel-artifacts/go/20250930-204932/{}"],
     version = "1.23.12fips",
 )

build/bazelutil/distdir_files.bzl

@@ -1188,13 +1188,13 @@ DISTDIR_FILES = {
     "https://storage.googleapis.com/public-bazel-artifacts/c-deps/20250801-193032/libproj_foreign.macos.20250801-193032.tar.gz": "8d28434cd175f0a32dfdd8ba8a5fa44c3d04d1e53cccfe9dbb3c6e301a03a47c",
     "https://storage.googleapis.com/public-bazel-artifacts/c-deps/20250801-193032/libproj_foreign.macosarm.20250801-193032.tar.gz": "a4b0bbb056bb462682b49ec34816f02c71047b38733d50d8de78b737c892db61",
     "https://storage.googleapis.com/public-bazel-artifacts/c-deps/20250801-193032/libproj_foreign.windows.20250801-193032.tar.gz": "a61f4faf7a7d017a194c64b453a38c982423ef3678fa049dbf114920759da59c",
-    "https://storage.googleapis.com/public-bazel-artifacts/go/20250818-202337/go1.23.12.darwin-amd64.tar.gz": "34457131f14281e21e25493d68e7519ccf26342d176dac36a4fc5dbf5ef122d9",
-    "https://storage.googleapis.com/public-bazel-artifacts/go/20250818-202337/go1.23.12.darwin-arm64.tar.gz": "30e0735ab9ccda203946536d24afe895abd1a1d3f35ad199f9768ccbdd5d60bc",
-    "https://storage.googleapis.com/public-bazel-artifacts/go/20250818-202337/go1.23.12.linux-amd64.tar.gz": "0cac0ac930ecb9458b8a0a7969cbf735c5884d24c879c97eb28a8997eca986fa",
-    "https://storage.googleapis.com/public-bazel-artifacts/go/20250818-202337/go1.23.12.linux-arm64.tar.gz": "528601fc8fb2c7e5ce8b7ae7651fd4fce2450bbef687beb96616edc5a9effa41",
-    "https://storage.googleapis.com/public-bazel-artifacts/go/20250818-202337/go1.23.12.linux-s390x.tar.gz": "f3f11bbb731da6716776d1c29a2db3d1063fa0a9f8c00636e6a77793ba79e2e3",
-    "https://storage.googleapis.com/public-bazel-artifacts/go/20250818-202337/go1.23.12.windows-amd64.tar.gz": "71b5b5b86b3a5ff9f124e21984abd874a6bfeb438f368de2eee7c60a25a19c94",
-    "https://storage.googleapis.com/public-bazel-artifacts/go/20250818-202337/go1.23.12fips.linux-amd64.tar.gz": "9c58fd7137b4c9d387a5c37fd2e728bc5d39357c7f8ba3358bcae513704c2983",
+    "https://storage.googleapis.com/public-bazel-artifacts/go/20250930-204932/go1.23.12.darwin-amd64.tar.gz": "25b853b77448c6f196475c1ab44d4a617b719db34e9e072c72c927a905bdbce6",
+    "https://storage.googleapis.com/public-bazel-artifacts/go/20250930-204932/go1.23.12.darwin-arm64.tar.gz": "fae1c6d45f72d559b9fe52224918321e7642cb75cd6a71186c7ccdf577189a4f",
+    "https://storage.googleapis.com/public-bazel-artifacts/go/20250930-204932/go1.23.12.linux-amd64.tar.gz": "03ac2c00dfde86a4beb3e94f063edffe1d2c089dfbbe1ec2776d1531c58e196b",
+    "https://storage.googleapis.com/public-bazel-artifacts/go/20250930-204932/go1.23.12.linux-arm64.tar.gz": "442f8d3510d141d6f479f05f7de00560aa1df59e0a5e1d02ca58ed9e4757d0d7",
+    "https://storage.googleapis.com/public-bazel-artifacts/go/20250930-204932/go1.23.12.linux-s390x.tar.gz": "6e3c09786b434bdd8d44de714656bef7c4fe484d2a2e6f8c98a797385e5f280e",
+    "https://storage.googleapis.com/public-bazel-artifacts/go/20250930-204932/go1.23.12.windows-amd64.tar.gz": "513af02afaa1c64501f1ec0a9c9445dbf88ea55c0b7a5a86e0b7fc94f0c59b39",
+    "https://storage.googleapis.com/public-bazel-artifacts/go/20250930-204932/go1.23.12fips.linux-amd64.tar.gz": "ae2d57fa43ef68aa70e6c0c1def065bd2806411acfc0d59591ddff81168be095",
     "https://storage.googleapis.com/public-bazel-artifacts/java/railroad/rr-1.63-java8.zip": "d2791cd7a44ea5be862f33f5a9b3d40aaad9858455828ebade7007ad7113fb41",
     "https://storage.googleapis.com/public-bazel-artifacts/js/rules_jest-v0.18.4.tar.gz": "d3bb833f74b8ad054e6bff5e41606ff10a62880cc99e4d480f4bdfa70add1ba7",
     "https://storage.googleapis.com/public-bazel-artifacts/js/rules_js-v1.42.3.tar.gz": "2cfb3875e1231cefd3fada6774f2c0c5a99db0070e0e48ea398acbff7c6c765b",

build/teamcity/internal/release/build-and-publish-patched-go/commit.txt

@@ -1 +1 @@
-309f11146b97839ffbba1ac245b7aa901e3dbbcb
+889a9764a479870035d3fb31d1910ce6c654533c

pkg/BUILD.bazel

@@ -1266,6 +1266,7 @@ GO_TARGETS = [
     "//pkg/cmd/roachprod/grafana:grafana",
     "//pkg/cmd/roachprod:roachprod",
     "//pkg/cmd/roachprod:roachprod_lib",
+    "//pkg/cmd/roachtest/cluster/mock:mockcluster",
     "//pkg/cmd/roachtest/cluster:cluster",
     "//pkg/cmd/roachtest/clusterstats:clusterstats",
     "//pkg/cmd/roachtest/clusterstats:clusterstats_test",

pkg/cmd/roachtest/cluster.go

@@ -3188,31 +3188,20 @@ func (c *clusterImpl) MaybeExtendCluster(
 // archForTest determines the CPU architecture to use for a test. If the test
 // doesn't specify it, one is chosen randomly depending on flags.
 func archForTest(ctx context.Context, l *logger.Logger, testSpec registry.TestSpec) vm.CPUArch {
-	if testSpec.Cluster.Arch != "" {
-		l.PrintfCtx(ctx, "Using specified arch=%q, %s", testSpec.Cluster.Arch, testSpec.Name)
-		return testSpec.Cluster.Arch
-	}
-
 	if roachtestflags.Cloud == spec.IBM {
 		// N.B. IBM only supports S390x on the "s390x" architecture.
 		l.PrintfCtx(ctx, "IBM Cloud: forcing arch=%q (only supported), %s", vm.ArchS390x, testSpec.Name)
 		return vm.ArchS390x
 	}
 
-	// CPU architecture is unspecified, choose one according to the
-	// probability distribution.
-	var arch vm.CPUArch
-	if prng.Float64() < roachtestflags.ARM64Probability {
-		arch = vm.ArchARM64
-	} else if prng.Float64() < roachtestflags.FIPSProbability {
-		// N.B. branch is taken with probability
-		//   (1 - arm64Probability) * fipsProbability
-		// which is P(fips & amd64).
-		// N.B. FIPS is only supported on 'amd64' at this time.
-		arch = vm.ArchFIPS
-	} else {
-		arch = vm.ArchAMD64
+	validArchs := spec.AllArchs
+	if !testSpec.Cluster.CompatibleArchs.IsEmpty() {
+		l.PrintfCtx(ctx, "Selecting from architectures=%q, %s", testSpec.Cluster.CompatibleArchs.String(), testSpec.Name)
+		validArchs = testSpec.Cluster.CompatibleArchs
 	}
+
+	arch := randomArch(ctx, l, validArchs, prng, roachtestflags.ARM64Probability, roachtestflags.FIPSProbability)
+
 	if roachtestflags.Cloud == spec.GCE && arch == vm.ArchARM64 {
 		// N.B. T2A support is rather limited, both in terms of supported
 		// regions and no local SSDs. Thus, we must fall back to AMD64 in
@@ -3232,6 +3221,61 @@ func archForTest(ctx context.Context, l *logger.Logger, testSpec registry.TestSp
 	return arch
 }
 
+// randomArch chooses a random architecture, respecting the set of valid architectures
+// specified by the test as well as the provided architecture probability flags.
+func randomArch(
+	ctx context.Context,
+	l *logger.Logger,
+	validArchs spec.ArchSet,
+	prng *rand.Rand,
+	arm64Probability, fipsProbability float64,
+) vm.CPUArch {
+	baseProbabilities := map[vm.CPUArch]float64{
+		vm.ArchAMD64: (1.0 - arm64Probability) * (1.0 - fipsProbability),
+		vm.ArchARM64: arm64Probability,
+		// N.B. FIPS is only supported on 'amd64' at this time:
+		// FIPS is taken with probability
+		//   (1 - arm64Probability) * fipsProbability
+		// 	 which is P(fips & amd64)
+		vm.ArchFIPS: (1.0 - arm64Probability) * fipsProbability,
+	}
+
+	// Calculate total weight for valid architectures only.
+	totalValidWeight := 0.0
+	validArchsList := validArchs.List()
+	for _, arch := range validArchsList {
+		totalValidWeight += baseProbabilities[arch]
+	}
+
+	// This would happen if the set of valid compatible arches (set by cluster spec) is disjoint with the set of
+	// enabled arches (set by roachtest flags).
+	if totalValidWeight == 0.0 {
+		l.PrintfCtx(ctx, "Defaulting to %s; CompatibleArches %s yields no architectures after applying roachtest arch probability flags", vm.ArchAMD64, validArchs.String())
+		return vm.ArchAMD64
+	}
+
+	// Since we allow only a subset of architectures, our total probability
+	// may not add up to 1. We normalize the weights amongst the valid architectures
+	// and track cumulative weights that give us "probability buckets" for each
+	// architecture.
+	cumulativeWeights := make([]float64, 0, len(validArchsList))
+	runningTotal := 0.0
+	for _, arch := range validArchsList {
+		normalizedWeight := baseProbabilities[arch] / totalValidWeight
+		runningTotal += normalizedWeight
+		cumulativeWeights = append(cumulativeWeights, runningTotal)
+	}
+	x := prng.Float64()
+	for i, weight := range cumulativeWeights {
+		if x < weight {
+			return validArchsList[i]
+		}
+	}
+	// Since we are adding floating point numbers, it's possible that we
+	// don't quite add up to 1.0. In that case, return the last architecture.
+	return validArchsList[len(validArchsList)-1]
+}
+
 // bucketVMsByProvider buckets cachedCluster.VMs by provider.
 func bucketVMsByProvider(cachedCluster *cloud.Cluster) map[string][]vm.VM {
 	providerToVMs := make(map[string][]vm.VM)

pkg/cmd/roachtest/cluster/mock/BUILD.bazel

@@ -0,0 +1,40 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "gomock")
+
+gomock(
+    name = "mock_cluster",
+    testonly = True,
+    out = "mock_cluster_generated.go",
+    interfaces = ["Cluster"],
+    library = "//pkg/cmd/roachtest/cluster",
+    package = "mockcluster",
+    self_package = "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster/mock",
+    visibility = [
+        "//pkg/cmd/roachtest/clusterstats:__pkg__",
+        "//pkg/cmd/roachtest/roachtestutil/mixedversion:__pkg__",
+        "//pkg/gen:__pkg__",
+    ],
+)
+
+go_library(
+    name = "mockcluster",
+    srcs = ["mock_cluster_generated.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster/mock",
+    visibility = [
+        "//pkg/cmd/roachtest/clusterstats:__pkg__",
+        "//pkg/cmd/roachtest/roachtestutil/mixedversion:__pkg__",
+        "//pkg/gen:__pkg__",
+    ],
+    deps = [
+        "//pkg/cmd/roachprod/grafana",
+        "//pkg/cmd/roachtest/cluster",
+        "//pkg/cmd/roachtest/option",
+        "//pkg/cmd/roachtest/spec",
+        "//pkg/roachprod",
+        "//pkg/roachprod/failureinjection/failures",
+        "//pkg/roachprod/install",
+        "//pkg/roachprod/logger",
+        "//pkg/roachprod/prometheus",
+        "//pkg/roachprod/vm",
+        "@com_github_golang_mock//gomock",
+    ],
+)

pkg/cmd/roachtest/clusterstats/mock_cluster_generated_test.go renamed to pkg/cmd/roachtest/cluster/mock/mock_cluster_generated.go

@@ -909,3 +909,119 @@ func TestVerifyLibraries(t *testing.T) {
 		})
 	}
 }
+
+func TestRandomArchProbabilities(t *testing.T) {
+	ctx := context.Background()
+
+	tests := []struct {
+		validArchs           spec.ArchSet
+		arm64Probability     float64
+		fipsProbability      float64
+		expectedDistribution map[vm.CPUArch]float64
+	}{
+		{
+			validArchs:       spec.AllArchs,
+			arm64Probability: 0.3,
+			fipsProbability:  0.2,
+			expectedDistribution: map[vm.CPUArch]float64{
+				vm.ArchAMD64: 0.56,
+				vm.ArchARM64: 0.3,
+				vm.ArchFIPS:  0.14,
+			},
+		},
+		{
+			validArchs:       spec.AllExceptFIPS,
+			arm64Probability: 0.4,
+			fipsProbability:  0.1,
+			expectedDistribution: map[vm.CPUArch]float64{
+				vm.ArchAMD64: 0.57447,
+				vm.ArchARM64: 0.42553,
+			},
+		},
+		{
+			validArchs:       spec.OnlyAMD64,
+			arm64Probability: 0.5,
+			fipsProbability:  0.3,
+			expectedDistribution: map[vm.CPUArch]float64{
+				vm.ArchAMD64: 1.0, // Only valid architecture
+			},
+		},
+		{
+			validArchs:       spec.OnlyARM64,
+			arm64Probability: 0.5,
+			fipsProbability:  0.3,
+			expectedDistribution: map[vm.CPUArch]float64{
+				vm.ArchARM64: 1.0, // Only valid architecture
+			},
+		},
+		{
+			validArchs:       spec.OnlyFIPS,
+			arm64Probability: 0.2,
+			fipsProbability:  0.0,
+			expectedDistribution: map[vm.CPUArch]float64{
+				vm.ArchAMD64: 1.0, // Should fall back to AMD64
+			},
+		},
+		{
+			validArchs:       spec.AllExceptFIPS,
+			arm64Probability: 0.0,
+			fipsProbability:  1.0,
+			expectedDistribution: map[vm.CPUArch]float64{
+				vm.ArchAMD64: 1.0, // Should fall back to AMD64
+			},
+		},
+		{
+			validArchs:       spec.AllExceptFIPS,
+			arm64Probability: 0.5,
+			fipsProbability:  1.0,
+			expectedDistribution: map[vm.CPUArch]float64{
+				vm.ArchARM64: 1.0,
+			},
+		},
+		{
+			validArchs:       spec.Archs(spec.ArchAMD64, spec.ArchFIPS),
+			arm64Probability: 0.3,
+			fipsProbability:  0.4,
+			expectedDistribution: map[vm.CPUArch]float64{
+				vm.ArchAMD64: 0.6,
+				vm.ArchFIPS:  0.4,
+			},
+		},
+		{
+			validArchs:       spec.Archs(spec.ArchARM64, spec.ArchFIPS),
+			arm64Probability: 0.6,
+			fipsProbability:  0.2,
+			expectedDistribution: map[vm.CPUArch]float64{
+				vm.ArchARM64: 0.88235,
+				vm.ArchFIPS:  0.11765,
+			},
+		},
+	}
+
+	// Since this is a statistical test, we want to use a fixed seed to avoid flakes,
+	// i.e. a 99% confidence interval would be expected to fail when stressed 100 times.
+	//
+	// We can run this manually with a random seed for more confidence in our distribution:
+	// prng, _ := randutil.NewTestRand()
+	prng := rand.New(rand.NewSource(12345))
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("%s/arm=%d%%/fips=%d%%", test.validArchs, int(test.arm64Probability*100), int(test.fipsProbability*100)), func(t *testing.T) {
+			const numSamples = 10000
+			counts := make(map[vm.CPUArch]int)
+
+			// Generate samples
+			for i := 0; i < numSamples; i++ {
+				arch := randomArch(ctx, nilLogger(), test.validArchs, prng, test.arm64Probability, test.fipsProbability)
+				counts[arch]++
+			}
+
+			for expectedArch, expectedProb := range test.expectedDistribution {
+				actualCount := float64(counts[expectedArch])
+				actualProb := actualCount / float64(numSamples)
+
+				tolerance := 0.02
+				require.InDelta(t, expectedProb, actualProb, tolerance)
+			}
+		})
+	}
+}

pkg/cmd/roachtest/cluster_test.go

@@ -15,7 +15,6 @@ go_library(
         "//pkg/cmd/roachprod-microbench/util",
         # Required for generated mocks
         "//pkg/cmd/roachprod/grafana",  #keep
-        "//pkg/cmd/roachtest/cluster",
         "//pkg/cmd/roachtest/option",
         "//pkg/cmd/roachtest/test",
         # Required for generated mocks
@@ -36,6 +35,7 @@ go_library(
         "@com_github_prometheus_client_golang//api/prometheus/v1:prometheus",
         "@com_github_prometheus_common//model",
         "//pkg/cmd/roachtest/roachtestutil",
+        "//pkg/cmd/roachtest/cluster",
     ],
 )
 
@@ -45,15 +45,15 @@ go_test(
         "exporter_test.go",
         "streamer_test.go",
         ":mock_client",  # keep
-        ":mock_cluster",  # keep
         ":mock_test",  # keep
     ],
     embed = [":clusterstats"],
     embedsrcs = ["openmetrics_expected.txt"],
     deps = [
         "//pkg/cmd/roachtest/cluster",
+        "//pkg/cmd/roachtest/cluster/mock:mockcluster",  #keep
         "//pkg/cmd/roachtest/registry",
-        # Required for generated mocks
+        "//pkg/cmd/roachtest/roachtestutil",
         "//pkg/cmd/roachtest/roachtestutil/task",  #keep
         "//pkg/cmd/roachtest/spec",
         "//pkg/cmd/roachtest/test",
@@ -64,7 +64,6 @@ go_test(
         "@com_github_prometheus_client_golang//api/prometheus/v1:prometheus",
         "@com_github_prometheus_common//model",
         "@com_github_stretchr_testify//require",
-        "//pkg/cmd/roachtest/roachtestutil",
     ],
 )
 
@@ -91,15 +90,3 @@ gomock(
         "//pkg/gen:__pkg__",
     ],
 )
-
-gomock(
-    name = "mock_cluster",
-    out = "mock_cluster_generated_test.go",
-    interfaces = ["Cluster"],
-    library = "//pkg/cmd/roachtest/cluster",
-    package = "clusterstats",
-    visibility = [
-        ":__pkg__",
-        "//pkg/gen:__pkg__",
-    ],
-)