@@ -826,8 +826,16 @@ func (t *Test) BackgroundCommand(
826
826
// if passed, is the command run to initialize the workload; it is run
827
827
// synchronously as a regular startup function. `runCmd` is the
828
828
// command to actually run the command; it is run in the background.
829
+ //
830
+ // If overrideBinary is true, the binary used to run the command(s) will
831
+ // be replaced with the cockroach binary of the current version the
832
+ // cluster is running in.
833
+ // TODO(testeng): Replace with https://github.com/cockroachdb/cockroach/issues/147374
829
834
func (t * Test ) Workload (
830
- name string , node option.NodeListOption , initCmd , runCmd * roachtestutil.Command ,
835
+ name string ,
836
+ node option.NodeListOption ,
837
+ initCmd , runCmd * roachtestutil.Command ,
838
+ overrideBinary bool ,
831
839
) StopFunc {
832
840
seed := uint64 (t .prng .Int63 ())
833
841
addSeed := func (cmd * roachtestutil.Command ) {
@@ -838,11 +846,31 @@ func (t *Test) Workload(
838
846
839
847
if initCmd != nil {
840
848
addSeed (initCmd )
841
- t .OnStartup (fmt .Sprintf ("initialize %s workload" , name ), t .runCommandFunc (node , initCmd .String ()))
849
+ t .OnStartup (fmt .Sprintf ("initialize %s workload" , name ), func (ctx context.Context , l * logger.Logger , rng * rand.Rand , h * Helper ) error {
850
+ if overrideBinary {
851
+ binary , err := clusterupgrade .UploadCockroach (ctx , t .rt , t .logger , t .cluster , node , h .System .FromVersion )
852
+ if err != nil {
853
+ t .rt .Fatal (err )
854
+ }
855
+ initCmd .Binary = binary
856
+ }
857
+ l .Printf ("running command `%s` on nodes %v" , initCmd .String (), node )
858
+ return t .cluster .RunE (ctx , option .WithNodes (node ), initCmd .String ())
859
+ })
842
860
}
843
861
844
862
addSeed (runCmd )
845
- return t .BackgroundCommand (fmt .Sprintf ("%s workload" , name ), node , runCmd )
863
+ return t .BackgroundFunc (fmt .Sprintf ("%s workload" , name ), func (ctx context.Context , l * logger.Logger , rng * rand.Rand , h * Helper ) error {
864
+ if overrideBinary {
865
+ binary , err := clusterupgrade .UploadCockroach (ctx , t .rt , t .logger , t .cluster , node , h .System .FromVersion )
866
+ if err != nil {
867
+ t .rt .Fatal (err )
868
+ }
869
+ runCmd .Binary = binary
870
+ }
871
+ l .Printf ("running command `%s` on nodes %v" , runCmd .String (), node )
872
+ return t .cluster .RunE (ctx , option .WithNodes (node ), runCmd .String ())
873
+ })
846
874
}
847
875
848
876
// Run is like RunE, except it fatals the test if any error occurs.
0 commit comments