dev: introduce --test-args for go test binaries

It's useful to be able to pass arguments down to the go test binary
directly. This commit introduces a top-level flag to do just that. Now
we can do the following sort of thing:

    dev bench pkg/bench -f='BenchmarkTracing/1node/scan/trace=off' \
      --test-args '-test.memprofile=mem.out -test.cpuprofile=cpu.out'

Release note: None

go.mod

@@ -163,6 +163,8 @@ require (
 	vitess.io/vitess v0.0.0-00010101000000-000000000000
 )
 
+require github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
+
 require (
 	cloud.google.com/go v0.100.2 // indirect
 	cloud.google.com/go/compute v1.1.0 // indirect

go.sum

@@ -1067,6 +1067,7 @@ github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1/go.mod h1:kpwsk12EmLe
 github.com/google/pprof v0.0.0-20210827144239-02619b876842 h1:JCrt5MIE1fHQtdy1825HwJ45oVQaqHE6lgssRhjcg/o=
 github.com/google/pprof v0.0.0-20210827144239-02619b876842/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
 github.com/google/skylark v0.0.0-20181101142754-a5f7082aabed h1:rZdD1GeRTHD1aG+VIvhQEYXurx6Wfg4QIT5YVl2tSC8=
 github.com/google/skylark v0.0.0-20181101142754-a5f7082aabed/go.mod h1:CKSX6SxHW1vp20ZNaeGe3TFFBIwCG6vaYrpAiOzX+NA=

pkg/cmd/dev/BUILD.bazel

@@ -28,6 +28,7 @@ go_library(
         "//pkg/cmd/dev/io/exec",
         "//pkg/cmd/dev/io/os",
         "@com_github_alessio_shellescape//:shellescape",
+        "@com_github_google_shlex//:shlex",
         "@com_github_spf13_cobra//:cobra",
     ],
 )
@@ -53,6 +54,7 @@ go_test(
         "//pkg/testutils",
         "@com_github_alessio_shellescape//:shellescape",
         "@com_github_cockroachdb_datadriven//:datadriven",
+        "@com_github_google_shlex//:shlex",
         "@com_github_irfansharif_recorder//:recorder",
         "@com_github_stretchr_testify//require",
     ],

pkg/cmd/dev/bench.go

@@ -46,6 +46,7 @@ func makeBenchCmd(runE func(cmd *cobra.Command, args []string) error) *cobra.Com
 	// `go help testflag`).
 	benchCmd.Flags().String(benchTimeFlag, "", "duration to run each benchmark for")
 	benchCmd.Flags().Bool(benchMemFlag, false, "print memory allocations for benchmarks")
+	benchCmd.Flags().String(testArgsFlag, "", "additional arguments to pass to go test binary")
 
 	return benchCmd
 }
@@ -63,6 +64,7 @@ func (d *dev) bench(cmd *cobra.Command, commandLine []string) error {
 		count       = mustGetFlagInt(cmd, countFlag)
 		benchTime   = mustGetFlagString(cmd, benchTimeFlag)
 		benchMem    = mustGetFlagBool(cmd, benchMemFlag)
+		testArgs    = mustGetFlagString(cmd, testArgsFlag)
 	)
 
 	// Enumerate all benches to run.
@@ -133,6 +135,13 @@ func (d *dev) bench(cmd *cobra.Command, commandLine []string) error {
 	if benchMem {
 		args = append(args, "--test_arg", "-test.benchmem")
 	}
+	if testArgs != "" {
+		goTestArgs, err := d.getGoTestArgs(ctx, testArgs)
+		if err != nil {
+			return err
+		}
+		args = append(args, goTestArgs...)
+	}
 	args = append(args, d.getTestOutputArgs(false /* stress */, verbose, showLogs)...)
 	args = append(args, additionalBazelArgs...)
 	logCommand("bazel", args...)

pkg/cmd/dev/datadriven_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cmd/dev/io/os"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/datadriven"
+	"github.com/google/shlex"
 	"github.com/stretchr/testify/require"
 )
 
@@ -90,19 +91,17 @@ func TestDataDriven(t *testing.T) {
 				dev.cli.SetOut(ioutil.Discard)
 			}
 
-			require.Equalf(t, d.Cmd, "dev", "unknown command: %s", d.Cmd)
-			var args []string
-			for _, cmdArg := range d.CmdArgs {
-				args = append(args, cmdArg.Key)
-				if len(cmdArg.Vals) != 0 {
-					args = append(args, cmdArg.Vals[0])
-				}
-			}
-			dev.cli.SetArgs(args)
+			require.Equalf(t, d.Cmd, "exec", "unknown command: %s", d.Cmd)
+			tokens, err := shlex.Split(d.Input)
+			require.NoError(t, err)
+			require.NotEmpty(t, tokens)
+			require.Equal(t, "dev", tokens[0])
+
+			dev.cli.SetArgs(tokens[1:])
+
 			if err := dev.cli.Execute(); err != nil {
 				return fmt.Sprintf("err: %s", err)
 			}
-
 			logs, err := ioutil.ReadAll(logger)
 			require.NoError(t, err)
 			return string(logs)

pkg/cmd/dev/io/exec/exec.go

@@ -249,13 +249,13 @@ func (e *Exec) commandContextImpl(
 // Next is a thin interceptor for all exec activity, running them through
 // testing knobs first.
 func (e *Exec) Next(command string, f func() (output string, err error)) (string, error) {
-	if e.knobs.dryrun {
-		return "", nil
-	}
 	if e.knobs.intercept != nil {
 		if output, ok := e.knobs.intercept[command]; ok {
 			return output, nil
 		}
 	}
+	if e.knobs.dryrun {
+		return "", nil
+	}
 	return e.Recorder.Next(command, f)
 }

pkg/cmd/dev/test.go

@@ -11,11 +11,13 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"path/filepath"
 	"strings"
 	"time"
 
+	"github.com/google/shlex"
 	"github.com/spf13/cobra"
 )
 
@@ -31,7 +33,7 @@ const (
 	raceFlag        = "race"
 	ignoreCacheFlag = "ignore-cache"
 	rewriteFlag     = "rewrite"
-	rewriteArgFlag  = "rewrite-arg"
+	testArgsFlag    = "test-args"
 	vModuleFlag     = "vmodule"
 )
 
@@ -69,7 +71,7 @@ func makeTestCmd(runE func(cmd *cobra.Command, args []string) error) *cobra.Comm
 	testCmd.Flags().Bool(raceFlag, false, "run tests using race builds")
 	testCmd.Flags().Bool(ignoreCacheFlag, false, "ignore cached test runs")
 	testCmd.Flags().String(rewriteFlag, "", "argument to pass to underlying test binary (only applicable to certain tests)")
-	testCmd.Flags().String(rewriteArgFlag, "", "additional argument to pass to -rewrite (implies --rewrite)")
+	testCmd.Flags().String(testArgsFlag, "", "additional arguments to pass to go test binary")
 	testCmd.Flags().Lookup(rewriteFlag).NoOptDefVal = "-rewrite"
 	testCmd.Flags().String(vModuleFlag, "", "comma-separated list of pattern=N settings for file-filtered logging")
 	return testCmd
@@ -83,7 +85,7 @@ func (d *dev) test(cmd *cobra.Command, commandLine []string) error {
 		ignoreCache   = mustGetFlagBool(cmd, ignoreCacheFlag)
 		race          = mustGetFlagBool(cmd, raceFlag)
 		rewrite       = mustGetFlagString(cmd, rewriteFlag)
-		rewriteArg    = mustGetFlagString(cmd, rewriteArgFlag)
+		testArgs      = mustGetFlagString(cmd, testArgsFlag)
 		short         = mustGetFlagBool(cmd, shortFlag)
 		stress        = mustGetFlagBool(cmd, stressFlag)
 		stressCmdArgs = mustGetFlagString(cmd, stressArgsFlag)
@@ -146,9 +148,6 @@ func (d *dev) test(cmd *cobra.Command, commandLine []string) error {
 		}
 		args = append(args, fmt.Sprintf("--test_env=COCKROACH_WORKSPACE=%s", workspace))
 		args = append(args, "--test_arg", rewrite)
-		if rewriteArg != "" {
-			args = append(args, "--test_arg", rewriteArg)
-		}
 		for _, testTarget := range testTargets {
 			dir := getDirectoryFromTarget(testTarget)
 			args = append(args, fmt.Sprintf("--sandbox_writable_path=%s", filepath.Join(workspace, dir)))
@@ -186,9 +185,13 @@ func (d *dev) test(cmd *cobra.Command, commandLine []string) error {
 	if vModule != "" {
 		args = append(args, "--test_arg", fmt.Sprintf("-vmodule=%s", vModule))
 	}
-	// TODO(irfansharif): Support --go-flags, to pass in an arbitrary set of
-	// flags into the go test binaries. Gives better coverage to everything
-	// listed under `go help testflags`.
+	if testArgs != "" {
+		goTestArgs, err := d.getGoTestArgs(ctx, testArgs)
+		if err != nil {
+			return err
+		}
+		args = append(args, goTestArgs...)
+	}
 
 	args = append(args, d.getTestOutputArgs(stress, verbose, showLogs)...)
 	args = append(args, additionalBazelArgs...)
@@ -203,6 +206,38 @@ func (d *dev) test(cmd *cobra.Command, commandLine []string) error {
 	//   [...] while parsing 'pkg/f:all': no such package 'pkg/f'
 }
 
+func (d *dev) getGoTestArgs(ctx context.Context, testArgs string) ([]string, error) {
+	var goTestArgs []string
+	if testArgs != "" {
+		// The test output directory defaults to wherever "go test" is running,
+		// which for us is somewhere in the sandbox. When passing arguments for
+		// go test (-{mem,cpu}profile, -trace for e.g.), we may be interested in
+		// output files. For that reason, configure it to write (if anything) to the
+		// workspace by default.
+		//
+		// TODO(irfansharif): We could also elevate the go test flags that do write
+		// output into top-level dev flags (like --rewrite) and only selectively
+		// configure -test.outputdir and the sandbox writable path.
+		workspace, err := d.getWorkspace(ctx)
+		if err != nil {
+			return nil, err
+		}
+		goTestArgs = append(goTestArgs,
+			"--test_arg", fmt.Sprintf("-test.outputdir=%s", workspace),
+			fmt.Sprintf("--sandbox_writable_path=%s", filepath.Join(workspace)),
+		)
+
+		parts, err := shlex.Split(testArgs)
+		if err != nil {
+			return nil, err
+		}
+		for _, part := range parts {
+			goTestArgs = append(goTestArgs, "--test_arg", part)
+		}
+	}
+	return goTestArgs, nil
+}
+
 func (d *dev) getTestOutputArgs(stress, verbose, showLogs bool) []string {
 	testOutputArgs := []string{"--test_output", "errors"}
 	if stress {

pkg/cmd/dev/testdata/datadriven/bench

@@ -1,15 +1,25 @@
+exec
 dev bench pkg/spanconfig/...
 ----
 bazel test pkg/spanconfig/...:all --test_arg -test.run=- --test_arg -test.bench=. --test_output errors
 
+exec
 dev bench pkg/sql/parser --filter=BenchmarkParse
 ----
 bazel test pkg/sql/parser:all --test_arg -test.run=- --test_arg -test.bench=BenchmarkParse --test_sharding_strategy=disabled --test_output errors
 
+exec
 dev bench pkg/bench -f=BenchmarkTracing/1node/scan/trace=off --count=2 --bench-time=10x --bench-mem
 ----
 bazel test pkg/bench:all --test_arg -test.run=- --test_arg -test.bench=BenchmarkTracing/1node/scan/trace=off --test_sharding_strategy=disabled --test_arg -test.count=2 --test_arg -test.benchtime=10x --test_arg -test.benchmem --test_output errors
 
+exec
 dev bench pkg/spanconfig/spanconfigkvsubscriber -f=BenchmarkSpanConfigDecoder --cpus=10 --ignore-cache -v --timeout=50s
 ----
 bazel test --local_cpu_resources=10 --test_timeout=50 pkg/spanconfig/spanconfigkvsubscriber:all --nocache_test_results --test_arg -test.run=- --test_arg -test.bench=BenchmarkSpanConfigDecoder --test_sharding_strategy=disabled --test_arg -test.v --test_output all
+
+exec
+dev bench pkg/bench -f='BenchmarkTracing/1node/scan/trace=off' --test-args '-test.memprofile=mem.out -test.cpuprofile=cpu.out'
+----
+bazel info workspace --color=no
+bazel test pkg/bench:all --test_arg -test.run=- --test_arg -test.bench=BenchmarkTracing/1node/scan/trace=off --test_sharding_strategy=disabled --test_arg -test.outputdir=crdb-checkout --sandbox_writable_path=crdb-checkout --test_arg -test.memprofile=mem.out --test_arg -test.cpuprofile=cpu.out --test_output errors

pkg/cmd/dev/testdata/datadriven/compose

@@ -1,7 +1,9 @@
+exec
 dev compose
 ----
 bazel run //pkg/compose:compose_test --config=test
 
+exec
 dev compose --cpus 12 --short --timeout 1m -f TestComposeCompare
 ----
 bazel run //pkg/compose:compose_test --config=test --local_cpu_resources=12 --test_filter=TestComposeCompare --test_arg -test.short --test_timeout=60

pkg/cmd/dev/testdata/datadriven/dev-build

@@ -1,62 +1,68 @@
+exec
 dev build cockroach-short --skip-generate
 ----
 bazel build //pkg/cmd/cockroach-short:cockroach-short
 bazel info workspace --color=no
-mkdir bin
+mkdir crdb-checkout/bin
 bazel info bazel-bin --color=no
-rm cockroach-short
-ln -s pkg/cmd/cockroach-short/cockroach-short_/cockroach-short cockroach-short
-rm cockroach
-ln -s pkg/cmd/cockroach-short/cockroach-short_/cockroach-short cockroach
+rm crdb-checkout/cockroach-short
+ln -s sandbox/pkg/cmd/cockroach-short/cockroach-short_/cockroach-short crdb-checkout/cockroach-short
+rm crdb-checkout/cockroach
+ln -s sandbox/pkg/cmd/cockroach-short/cockroach-short_/cockroach-short crdb-checkout/cockroach
 
+exec
 dev build cockroach-short --cpus=12 --skip-generate
 ----
 bazel build --local_cpu_resources=12 //pkg/cmd/cockroach-short:cockroach-short
 bazel info workspace --color=no
-mkdir bin
+mkdir crdb-checkout/bin
 bazel info bazel-bin --color=no
-rm cockroach-short
-ln -s pkg/cmd/cockroach-short/cockroach-short_/cockroach-short cockroach-short
-rm cockroach
-ln -s pkg/cmd/cockroach-short/cockroach-short_/cockroach-short cockroach
+rm crdb-checkout/cockroach-short
+ln -s sandbox/pkg/cmd/cockroach-short/cockroach-short_/cockroach-short crdb-checkout/cockroach-short
+rm crdb-checkout/cockroach
+ln -s sandbox/pkg/cmd/cockroach-short/cockroach-short_/cockroach-short crdb-checkout/cockroach
 
+exec
 dev build --debug short --skip-generate
 ----
 bazel build //pkg/cmd/cockroach-short:cockroach-short
 bazel info workspace --color=no
-mkdir bin
+mkdir crdb-checkout/bin
 bazel info bazel-bin --color=no
-rm cockroach-short
-ln -s pkg/cmd/cockroach-short/cockroach-short_/cockroach-short cockroach-short
-rm cockroach
-ln -s pkg/cmd/cockroach-short/cockroach-short_/cockroach-short cockroach
+rm crdb-checkout/cockroach-short
+ln -s sandbox/pkg/cmd/cockroach-short/cockroach-short_/cockroach-short crdb-checkout/cockroach-short
+rm crdb-checkout/cockroach
+ln -s sandbox/pkg/cmd/cockroach-short/cockroach-short_/cockroach-short crdb-checkout/cockroach
 
+exec
 dev build short --skip-generate -- -s
 ----
 bazel build //pkg/cmd/cockroach-short:cockroach-short -s
 bazel info workspace --color=no
-mkdir bin
+mkdir crdb-checkout/bin
 bazel info bazel-bin --color=no
-rm cockroach-short
-ln -s pkg/cmd/cockroach-short/cockroach-short_/cockroach-short cockroach-short
-rm cockroach
-ln -s pkg/cmd/cockroach-short/cockroach-short_/cockroach-short cockroach
+rm crdb-checkout/cockroach-short
+ln -s sandbox/pkg/cmd/cockroach-short/cockroach-short_/cockroach-short crdb-checkout/cockroach-short
+rm crdb-checkout/cockroach
+ln -s sandbox/pkg/cmd/cockroach-short/cockroach-short_/cockroach-short crdb-checkout/cockroach
 
+exec
 dev build --skip-generate -- --verbose_failures --sandbox_debug
 ----
 bazel run @nodejs//:yarn -- --check-files --cwd pkg/ui --offline
 bazel build //pkg/cmd/cockroach:cockroach --config=with_ui --verbose_failures --sandbox_debug
 bazel info workspace --color=no
-mkdir bin
+mkdir crdb-checkout/bin
 bazel info bazel-bin --color=no
-rm cockroach
-ln -s pkg/cmd/cockroach/cockroach_/cockroach cockroach
+rm crdb-checkout/cockroach
+ln -s sandbox/pkg/cmd/cockroach/cockroach_/cockroach crdb-checkout/cockroach
 
+exec
 dev build stress --skip-generate
 ----
 bazel build @com_github_cockroachdb_stress//:stress
 bazel info workspace --color=no
-mkdir bin
+mkdir crdb-checkout/bin
 bazel info bazel-bin --color=no
-rm bin/stress
-ln -s external/com_github_cockroachdb_stress/stress_/stress bin/stress
+rm crdb-checkout/bin/stress
+ln -s sandbox/external/com_github_cockroachdb_stress/stress_/stress crdb-checkout/bin/stress

pkg/cmd/dev/testdata/datadriven/generate

@@ -1,15 +1,18 @@
+exec
 dev gen protobuf
 ----
 bazel run //pkg/gen:go_proto
 
+exec
 dev gen bazel
 ----
 bazel info workspace --color=no
-build/bazelutil/bazel-generate.sh
+crdb-checkout/build/bazelutil/bazel-generate.sh
 
+exec
 dev generate bazel --mirror --force
 ----
 bazel info workspace --color=no
 export COCKROACH_BAZEL_CAN_MIRROR=1
 export COCKROACH_BAZEL_FORCE_GENERATE=1
-build/bazelutil/bazel-generate.sh
+crdb-checkout/build/bazelutil/bazel-generate.sh