roachtest: add --cockroach-stage flag for remote binary staging

Add --cockroach-stage flag that stages binaries directly from cloud storage
instead of uploading local files. The flag accepts version/SHA values like
"latest", "v23.2.0", or commit SHAs, is mutually exclusive with --cockroach,
and normalizes "latest" to empty string for proper .LATEST URL construction.

Release note: none.
Epic: none

Author: dt

pkg/cmd/roachtest/cluster.go

@@ -1932,9 +1932,15 @@ func (c *clusterImpl) PutE(
 }
 
 // PutCockroach uploads a binary with or without runtime assertions enabled,
-// as determined by t.Cockroach(). Note that we upload to all nodes even if they
-// don't use the binary, so that the test runner can always fetch logs.
+// as determined by t.Cockroach(). If --cockroach-stage flag is set, it stages
+// the binary from cloud storage instead of uploading a local binary.
+// Note that we upload/stage to all nodes even if they don't use the binary,
+// so that the test runner can always fetch logs.
 func (c *clusterImpl) PutCockroach(ctx context.Context, l *logger.Logger, t *testImpl) error {
+	if roachtestflags.CockroachStage != "" {
+		// Use staging instead of upload when --cockroach-stage is specified
+		return c.Stage(ctx, l, "cockroach", roachtestflags.CockroachStage, ".", c.All())
+	}
 	return c.PutE(ctx, l, t.Cockroach(), test.DefaultCockroachPath, c.All())
 }
 

pkg/cmd/roachtest/main.go

@@ -573,4 +573,16 @@ func validateAndConfigure(cmd *cobra.Command, args []string) {
 	if roachtestflags.SelectiveTests && selectProbFlagInfo != nil {
 		printErrAndExit(fmt.Errorf("select-probability and selective-tests=true are incompatible. Disable one of them"))
 	}
+
+	// --cockroach and --cockroach-stage flags are mutually exclusive.
+	cockroachFlagInfo := roachtestflags.Changed(&roachtestflags.CockroachPath)
+	cockroachStageFlagInfo := roachtestflags.Changed(&roachtestflags.CockroachStage)
+	if cockroachFlagInfo != nil && cockroachStageFlagInfo != nil {
+		printErrAndExit(fmt.Errorf("--cockroach and --cockroach-stage are mutually exclusive. Use one or the other"))
+	}
+
+	// Normalize "latest" to empty string for staging system
+	if roachtestflags.CockroachStage == "latest" {
+		roachtestflags.CockroachStage = ""
+	}
 }

pkg/cmd/roachtest/roachtestflags/flags.go

@@ -112,6 +112,15 @@ var (
 		Usage: `Absolute path to cockroach binary to use`,
 	})
 
+	CockroachStage string
+	_              = registerRunFlag(&CockroachStage, FlagInfo{
+		Name: "cockroach-stage",
+		Usage: `
+			Stage cockroach binary from cloud storage instead of uploading local binary.
+			Specify version/SHA (e.g., "latest", "v23.2.0", or commit SHA).
+			Mutually exclusive with --cockroach.`,
+	})
+
 	ConfigPath string
 	_          = registerRunOpsFlag(&ConfigPath, FlagInfo{
 		Name: "config",