Support partitions for other tests.

PiperOrigin-RevId: 345399936

16 files changed, 112 insertions, 122 deletions

diff --git a/Makefile b/Makefile
index 47d89c438..3f4c534ff 100644
--- a/Makefile
+++ b/Makefile
@@ -123,6 +123,7 @@ list-images: ## List all available images.
 ##
 PARTITION        ?= 1
 TOTAL_PARTITIONS ?= 1
+PARTITIONS       := --test_arg=--partition=$(PARTITION) --test_arg=--total_partitions=$(TOTAL_PARTITIONS)
 
 runsc: ## Builds the runsc binary.
 	@$(call submake,build OPTIONS="-c opt" TARGETS="//runsc")
@@ -137,7 +138,7 @@ smoke-tests: ## Runs a simple smoke test after build runsc.
 .PHONY: smoke-tests
 
 fuse-tests:
-	@$(call submake,test OPTIONS="--test_tag_filters fuse" TARGETS="test/fuse/...")
+	@$(call submake,test OPTIONS="--test_tag_filters fuse $(PARTITIONS)" TARGETS="test/fuse/...")
 .PHONY: fuse-tests
 
 unit-tests: ## Local package unit tests in pkg/..., runsc/, tools/.., etc.
@@ -161,28 +162,22 @@ network-tests: iptables-tests packetdrill-tests packetimpact-tests
 INTEGRATION_TARGETS := //test/image:image_test //test/e2e:integration_test
 
 syscall-%-tests:
-	@$(call submake,test OPTIONS="--test_tag_filters runsc_$*" TARGETS="test/syscalls/...")
+	@$(call submake,test OPTIONS="--test_tag_filters runsc_$* $(PARTITIONS)" TARGETS="test/syscalls/...")
 
 syscall-native-tests:
-	@$(call submake,test OPTIONS="--test_tag_filters native" TARGETS="test/syscalls/...")
+	@$(call submake,test OPTIONS="--test_tag_filters native $(PARTITIONS)" TARGETS="test/syscalls/...")
 .PHONY: syscall-native-tests
 
 syscall-tests: ## Run all system call tests.
-	@$(call submake,test TARGETS="test/syscalls/...")
+	@$(call submake,test OPTIONS="$(PARTITIONS)" TARGETS="test/syscalls/...")
 
 %-runtime-tests: load-runtimes_%
 	@$(call submake,install-runtime)
-	@$(call submake,test-runtime OPTIONS="--test_timeout=10800 --test_arg=--partition=$(PARTITION) --test_arg=--total_partitions=$(TOTAL_PARTITIONS)" TARGETS="//test/runtimes:$*")
+	@$(call submake,test-runtime OPTIONS="--test_timeout=10800" TARGETS="//test/runtimes:$*")
 
 %-runtime-tests_vfs2: load-runtimes_%
-ifeq ($(PARTITION),)
-	@$(eval PARTITION := 1)
-endif
-ifeq ($(TOTAL_PARTITIONS),)
-	@$(eval TOTAL_PARTITIONS := 1)
-endif
 	@$(call submake,install-runtime RUNTIME="vfs2" ARGS="--vfs2")
-	@$(call submake,test-runtime RUNTIME="vfs2" OPTIONS="--test_timeout=10800 --test_arg=--partition=$(PARTITION) --test_arg=--total_partitions=$(TOTAL_PARTITIONS)" TARGETS="//test/runtimes:$*")
+	@$(call submake,test-runtime RUNTIME="vfs2" OPTIONS="--test_timeout=10800" TARGETS="//test/runtimes:$*")
 
 do-tests: runsc
 	@$(call submake,run TARGETS="//runsc" ARGS="--rootless do true")
@@ -464,7 +459,7 @@ configure: ## Configures a single runtime. Requires sudo. Typically called from
 .PHONY: configure
 
 test-runtime: ## A convenient wrapper around test that provides the runtime argument. Target must still be provided.
-	@$(call submake,test OPTIONS="$(OPTIONS) --test_arg=--runtime=$(RUNTIME)")
+	@$(call submake,test OPTIONS="$(OPTIONS) --test_arg=--runtime=$(RUNTIME) $(PARTITIONS)")
 .PHONY: test-runtime
 
 nogo: ## Surfaces all nogo findings.
diff --git a/pkg/tcpip/stack/BUILD b/pkg/tcpip/stack/BUILD
index d87b809fd..29c4668b7 100644
--- a/pkg/tcpip/stack/BUILD
+++ b/pkg/tcpip/stack/BUILD
@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "go_library", "go_test")
+load("//tools:defs.bzl", "go_library", "go_test", "most_shards")
 load("//tools/go_generics:defs.bzl", "go_template_instance")
 
 package(licenses = ["notice"])
@@ -112,7 +112,7 @@ go_test(
         "transport_demuxer_test.go",
         "transport_test.go",
     ],
-    shard_count = 20,
+    shard_count = most_shards,
     deps = [
         ":stack",
         "//pkg/rand",
diff --git a/pkg/tcpip/transport/tcp/BUILD b/pkg/tcpip/transport/tcp/BUILD
index 3d8174a4f..cf232b508 100644
--- a/pkg/tcpip/transport/tcp/BUILD
+++ b/pkg/tcpip/transport/tcp/BUILD
@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "go_library", "go_test")
+load("//tools:defs.bzl", "go_library", "go_test", "more_shards")
 load("//tools/go_generics:defs.bzl", "go_template_instance")
 
 package(licenses = ["notice"])
@@ -93,7 +93,7 @@ go_test(
         "tcp_test.go",
         "tcp_timestamp_test.go",
     ],
-    shard_count = 10,
+    shard_count = more_shards,
     deps = [
         ":tcp",
         "//pkg/rand",
diff --git a/pkg/test/testutil/testutil.go b/pkg/test/testutil/testutil.go
index 757a689f6..4e115718d 100644
--- a/pkg/test/testutil/testutil.go
+++ b/pkg/test/testutil/testutil.go
@@ -48,7 +48,9 @@ import (
 )
 
 var (
-	checkpoint = flag.Bool("checkpoint", true, "control checkpoint/restore support")
+	checkpoint      = flag.Bool("checkpoint", true, "control checkpoint/restore support")
+	partition       = flag.Int("partition", 1, "partition number, this is 1-indexed")
+	totalPartitions = flag.Int("total_partitions", 1, "total number of partitions")
 )
 
 // IsCheckpointSupported returns the relevant command line flag.
@@ -521,7 +523,8 @@ func TouchShardStatusFile() error {
 }
 
 // TestIndicesForShard returns indices for this test shard based on the
-// TEST_SHARD_INDEX and TEST_TOTAL_SHARDS environment vars.
+// TEST_SHARD_INDEX and TEST_TOTAL_SHARDS environment vars, as well as
+// the passed partition flags.
 //
 // If either of the env vars are not present, then the function will return all
 // tests. If there are more shards than there are tests, then the returned list
@@ -546,6 +549,11 @@ func TestIndicesForShard(numTests int) ([]int, error) {
 		}
 	}
 
+	// Combine with the partitions.
+	partitionSize := shardTotal
+	shardTotal = (*totalPartitions) * shardTotal
+	shardIndex = partitionSize*(*partition-1) + shardIndex
+
 	// Calculate!
 	var indices []int
 	numBlocks := int(math.Ceil(float64(numTests) / float64(shardTotal)))
diff --git a/runsc/container/BUILD b/runsc/container/BUILD
index 1900ecceb..8793c8916 100644
--- a/runsc/container/BUILD
+++ b/runsc/container/BUILD
@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "go_library", "go_test")
+load("//tools:defs.bzl", "go_library", "go_test", "more_shards")
 
 package(licenses = ["notice"])
 
@@ -49,7 +49,7 @@ go_test(
         "//test/cmd/test_app",
     ],
     library = ":container",
-    shard_count = 10,
+    shard_count = more_shards,
     tags = [
         "requires-kvm",
     ],
diff --git a/test/packetdrill/defs.bzl b/test/packetdrill/defs.bzl
index fc28ce9ba..a6cbcc376 100644
--- a/test/packetdrill/defs.bzl
+++ b/test/packetdrill/defs.bzl
@@ -15,7 +15,7 @@ def _packetdrill_test_impl(ctx):
         # Make sure that everything is readable here.
         "find . -type f -exec chmod a+rx {} \\;",
         "find . -type d -exec chmod a+rx {} \\;",
-        "%s %s --init_script %s $@ -- %s\n" % (
+        "%s %s --init_script %s \"$@\" -- %s\n" % (
             test_runner.short_path,
             " ".join(ctx.attr.flags),
             ctx.files._init_script[0].short_path,
@@ -80,9 +80,7 @@ def packetdrill_netstack_test(name, **kwargs):
         kwargs["tags"] = PACKETDRILL_TAGS
     _packetdrill_test(
         name = name,
-        # This is the default runtime unless
-        # "--test_arg=--runtime=OTHER_RUNTIME" is used to override the value.
-        flags = ["--dut_platform", "netstack", "--runtime", "runsc-d"],
+        flags = ["--dut_platform", "netstack"],
         **kwargs
     )
 
diff --git a/test/packetdrill/packetdrill_test.sh b/test/packetdrill/packetdrill_test.sh
index 922547d65..d25cad83a 100755
--- a/test/packetdrill/packetdrill_test.sh
+++ b/test/packetdrill/packetdrill_test.sh
@@ -29,7 +29,7 @@ function failure() {
 }
 trap 'failure ${LINENO} "$BASH_COMMAND"' ERR
 
-declare -r LONGOPTS="dut_platform:,init_script:,runtime:"
+declare -r LONGOPTS="dut_platform:,init_script:,runtime:,partition:,total_partitions:"
 
 # Don't use declare below so that the error from getopt will end the script.
 PARSED=$(getopt --options "" --longoptions=$LONGOPTS --name "$0" -- "$@")
@@ -48,12 +48,17 @@ while true; do
       shift 2
       ;;
     --runtime)
-      # Not readonly because there might be multiple --runtime arguments and we
-      # want to use just the last one.  Only used if --dut_platform is
-      # "netstack".
       declare RUNTIME="$2"
       shift 2
       ;;
+    --partition)
+      # Ignored.
+      shift 2
+      ;;
+    --total_partitions)
+      # Ignored.
+      shift 2
+      ;;
     --)
       shift
       break
diff --git a/test/perf/BUILD b/test/perf/BUILD
index b763be50e..e25f090ae 100644
--- a/test/perf/BUILD
+++ b/test/perf/BUILD
@@ -1,3 +1,4 @@
+load("//tools:defs.bzl", "more_shards")
 load("//test/runner:defs.bzl", "syscall_test")
 
 package(licenses = ["notice"])
@@ -37,7 +38,7 @@ syscall_test(
 syscall_test(
     size = "enormous",
     debug = False,
-    shard_count = 10,
+    shard_count = more_shards,
     tags = ["nogotsan"],
     test = "//test/perf/linux:getdents_benchmark",
 )
diff --git a/test/runner/defs.bzl b/test/runner/defs.bzl
index 7618f6a21..829247657 100644
--- a/test/runner/defs.bzl
+++ b/test/runner/defs.bzl
@@ -12,7 +12,7 @@ def _runner_test_impl(ctx):
         "  mkdir -p \"${TEST_UNDECLARED_OUTPUTS_DIR}\"",
         "  chmod a+rwx \"${TEST_UNDECLARED_OUTPUTS_DIR}\"",
         "fi",
-        "exec %s %s %s\n" % (
+        "exec %s %s \"$@\" %s\n" % (
             ctx.files.runner[0].short_path,
             " ".join(ctx.attr.runner_args),
             ctx.files.test[0].short_path,
@@ -52,8 +52,6 @@ _runner_test = rule(
 
 def _syscall_test(
         test,
-        shard_count,
-        size,
         platform,
         use_tmpfs,
         tags,
@@ -63,7 +61,8 @@ def _syscall_test(
         overlay = False,
         add_uds_tree = False,
         vfs2 = False,
-        fuse = False):
+        fuse = False,
+        **kwargs):
     # Prepend "runsc" to non-native platform names.
     full_platform = platform if platform == "native" else "runsc_" + platform
 
@@ -126,15 +125,12 @@ def _syscall_test(
         name = name,
         test = test,
         runner_args = runner_args,
-        size = size,
         tags = tags,
-        shard_count = shard_count,
+        **kwargs
     )
 
 def syscall_test(
         test,
-        shard_count = 5,
-        size = "small",
         use_tmpfs = False,
         add_overlay = False,
         add_uds_tree = False,
@@ -142,18 +138,21 @@ def syscall_test(
         vfs2 = True,
         fuse = False,
         debug = True,
-        tags = None):
+        tags = None,
+        **kwargs):
     """syscall_test is a macro that will create targets for all platforms.
 
     Args:
       test: the test target.
-      shard_count: shards for defined tests.
-      size: the defined test size.
       use_tmpfs: use tmpfs in the defined tests.
       add_overlay: add an overlay test.
       add_uds_tree: add a UDS test.
       add_hostinet: add a hostinet test.
+      vfs2: enable VFS2 support.
+      fuse: enable FUSE support.
+      debug: enable debug output.
       tags: starting test tags.
+      **kwargs: additional test arguments.
     """
     if not tags:
         tags = []
@@ -173,8 +172,6 @@ def syscall_test(
 
     _syscall_test(
         test = test,
-        shard_count = shard_count,
-        size = size,
         platform = default_platform,
         use_tmpfs = use_tmpfs,
         add_uds_tree = add_uds_tree,
@@ -182,6 +179,7 @@ def syscall_test(
         debug = debug,
         vfs2 = True,
         fuse = fuse,
+        **kwargs
     )
     if fuse:
         # Only generate *_vfs2_fuse target if fuse parameter is enabled.
@@ -189,38 +187,35 @@ def syscall_test(
 
     _syscall_test(
         test = test,
-        shard_count = shard_count,
-        size = size,
         platform = "native",
         use_tmpfs = False,
         add_uds_tree = add_uds_tree,
         tags = list(tags),
         debug = debug,
+        **kwargs
     )
 
     for (platform, platform_tags) in platforms.items():
         _syscall_test(
             test = test,
-            shard_count = shard_count,
-            size = size,
             platform = platform,
             use_tmpfs = use_tmpfs,
             add_uds_tree = add_uds_tree,
             tags = platform_tags + tags,
             debug = debug,
+            **kwargs
         )
 
     if add_overlay:
         _syscall_test(
             test = test,
-            shard_count = shard_count,
-            size = size,
             platform = default_platform,
             use_tmpfs = use_tmpfs,
             add_uds_tree = add_uds_tree,
             tags = platforms[default_platform] + tags,
             debug = debug,
             overlay = True,
+            **kwargs
         )
 
         # TODO(gvisor.dev/issue/4407): Remove tags to enable VFS2 overlay tests.
@@ -230,8 +225,6 @@ def syscall_test(
         overlay_vfs2_tags.append("notap")
         _syscall_test(
             test = test,
-            shard_count = shard_count,
-            size = size,
             platform = default_platform,
             use_tmpfs = use_tmpfs,
             add_uds_tree = add_uds_tree,
@@ -239,38 +232,35 @@ def syscall_test(
             debug = debug,
             overlay = True,
             vfs2 = True,
+            **kwargs
         )
 
     if add_hostinet:
         _syscall_test(
             test = test,
-            shard_count = shard_count,
-            size = size,
             platform = default_platform,
             use_tmpfs = use_tmpfs,
             network = "host",
             add_uds_tree = add_uds_tree,
             tags = platforms[default_platform] + tags,
             debug = debug,
+            **kwargs
         )
 
     if not use_tmpfs:
         # Also test shared gofer access.
         _syscall_test(
             test = test,
-            shard_count = shard_count,
-            size = size,
             platform = default_platform,
             use_tmpfs = use_tmpfs,
             add_uds_tree = add_uds_tree,
             tags = platforms[default_platform] + tags,
             debug = debug,
             file_access = "shared",
+            **kwargs
         )
         _syscall_test(
             test = test,
-            shard_count = shard_count,
-            size = size,
             platform = default_platform,
             use_tmpfs = use_tmpfs,
             add_uds_tree = add_uds_tree,
@@ -278,4 +268,5 @@ def syscall_test(
             debug = debug,
             file_access = "shared",
             vfs2 = True,
+            **kwargs
         )
diff --git a/test/runtimes/BUILD b/test/runtimes/BUILD
index 22b526f59..510ffe013 100644
--- a/test/runtimes/BUILD
+++ b/test/runtimes/BUILD
@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "bzl_library")
+load("//tools:defs.bzl", "bzl_library", "more_shards", "most_shards")
 load("//test/runtimes:defs.bzl", "runtime_test")
 
 package(licenses = ["notice"])
@@ -7,7 +7,7 @@ runtime_test(
     name = "go1.12",
     exclude_file = "exclude/go1.12.csv",
     lang = "go",
-    shard_count = 8,
+    shard_count = more_shards,
 )
 
 runtime_test(
@@ -15,28 +15,28 @@ runtime_test(
     batch = 100,
     exclude_file = "exclude/java11.csv",
     lang = "java",
-    shard_count = 16,
+    shard_count = most_shards,
 )
 
 runtime_test(
     name = "nodejs12.4.0",
     exclude_file = "exclude/nodejs12.4.0.csv",
     lang = "nodejs",
-    shard_count = 8,
+    shard_count = most_shards,
 )
 
 runtime_test(
     name = "php7.3.6",
     exclude_file = "exclude/php7.3.6.csv",
     lang = "php",
-    shard_count = 8,
+    shard_count = more_shards,
 )
 
 runtime_test(
     name = "python3.7.3",
     exclude_file = "exclude/python3.7.3.csv",
     lang = "python",
-    shard_count = 8,
+    shard_count = more_shards,
 )
 
 bzl_library(
diff --git a/test/runtimes/runner/lib/lib.go b/test/runtimes/runner/lib/lib.go
index 2f1083baf..9272137ff 100644
--- a/test/runtimes/runner/lib/lib.go
+++ b/test/runtimes/runner/lib/lib.go
@@ -34,12 +34,7 @@ import (
 // RunTests is a helper that is called by main. It exists so that we can run
 // defered functions before exiting. It returns an exit code that should be
 // passed to os.Exit.
-func RunTests(lang, image, excludeFile string, partitionNum, totalPartitions, batchSize int, timeout time.Duration) int {
-	if partitionNum <= 0 || totalPartitions <= 0 || partitionNum > totalPartitions {
-		fmt.Fprintf(os.Stderr, "invalid partition %d of %d", partitionNum, totalPartitions)
-		return 1
-	}
-
+func RunTests(lang, image, excludeFile string, batchSize int, timeout time.Duration) int {
 	// TODO(gvisor.dev/issue/1624): Remove those tests from all exclude lists
 	// that only fail with VFS1.
 
@@ -63,7 +58,7 @@ func RunTests(lang, image, excludeFile string, partitionNum, totalPartitions, ba
 	// Get a slice of tests to run. This will also start a single Docker
 	// container that will be used to run each test. The final test will
 	// stop the Docker container.
-	tests, err := getTests(ctx, d, lang, image, partitionNum, totalPartitions, batchSize, timeout, excludes)
+	tests, err := getTests(ctx, d, lang, image, batchSize, timeout, excludes)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "%s\n", err.Error())
 		return 1
@@ -74,7 +69,7 @@ func RunTests(lang, image, excludeFile string, partitionNum, totalPartitions, ba
 }
 
 // getTests executes all tests as table tests.
-func getTests(ctx context.Context, d *dockerutil.Container, lang, image string, partitionNum, totalPartitions, batchSize int, timeout time.Duration, excludes map[string]struct{}) ([]testing.InternalTest, error) {
+func getTests(ctx context.Context, d *dockerutil.Container, lang, image string, batchSize int, timeout time.Duration, excludes map[string]struct{}) ([]testing.InternalTest, error) {
 	// Start the container.
 	opts := dockerutil.RunOpts{
 		Image: fmt.Sprintf("runtimes/%s", image),
@@ -90,18 +85,9 @@ func getTests(ctx context.Context, d *dockerutil.Container, lang, image string,
 		return nil, fmt.Errorf("docker exec failed: %v", err)
 	}
 
-	// Calculate a subset of tests to run corresponding to the current
-	// shard.
+	// Calculate a subset of tests.
 	tests := strings.Fields(list)
 	sort.Strings(tests)
-
-	partitionSize := len(tests) / totalPartitions
-	if partitionNum == totalPartitions {
-		tests = tests[(partitionNum-1)*partitionSize:]
-	} else {
-		tests = tests[(partitionNum-1)*partitionSize : partitionNum*partitionSize]
-	}
-
 	indices, err := testutil.TestIndicesForShard(len(tests))
 	if err != nil {
 		return nil, fmt.Errorf("TestsForShard() failed: %v", err)
diff --git a/test/runtimes/runner/main.go b/test/runtimes/runner/main.go
index 5b3443e36..ec79a22c2 100644
--- a/test/runtimes/runner/main.go
+++ b/test/runtimes/runner/main.go
@@ -25,13 +25,11 @@ import (
 )
 
 var (
-	lang            = flag.String("lang", "", "language runtime to test")
-	image           = flag.String("image", "", "docker image with runtime tests")
-	excludeFile     = flag.String("exclude_file", "", "file containing list of tests to exclude, in CSV format with fields: test name, bug id, comment")
-	partition       = flag.Int("partition", 1, "partition number, this is 1-indexed")
-	totalPartitions = flag.Int("total_partitions", 1, "total number of partitions")
-	batchSize       = flag.Int("batch", 50, "number of test cases run in one command")
-	timeout         = flag.Duration("timeout", 90*time.Minute, "batch timeout")
+	lang        = flag.String("lang", "", "language runtime to test")
+	image       = flag.String("image", "", "docker image with runtime tests")
+	excludeFile = flag.String("exclude_file", "", "file containing list of tests to exclude, in CSV format with fields: test name, bug id, comment")
+	batchSize   = flag.Int("batch", 50, "number of test cases run in one command")
+	timeout     = flag.Duration("timeout", 90*time.Minute, "batch timeout")
 )
 
 func main() {
@@ -40,5 +38,5 @@ func main() {
 		fmt.Fprintf(os.Stderr, "lang and image flags must not be empty\n")
 		os.Exit(1)
 	}
-	os.Exit(lib.RunTests(*lang, *image, *excludeFile, *partition, *totalPartitions, *batchSize, *timeout))
+	os.Exit(lib.RunTests(*lang, *image, *excludeFile, *batchSize, *timeout))
 }
diff --git a/test/syscalls/BUILD b/test/syscalls/BUILD
index b5a4ef4df..135d58ae6 100644
--- a/test/syscalls/BUILD
+++ b/test/syscalls/BUILD
@@ -1,3 +1,4 @@
+load("//tools:defs.bzl", "more_shards", "most_shards")
 load("//test/runner:defs.bzl", "syscall_test")
 
 package(licenses = ["notice"])
@@ -12,7 +13,7 @@ syscall_test(
 
 syscall_test(
     size = "large",
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:accept_bind_test",
 )
 
@@ -32,7 +33,7 @@ syscall_test(
 
 syscall_test(
     size = "medium",
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:alarm_test",
 )
 
@@ -66,7 +67,7 @@ syscall_test(
     size = "large",
     # Produce too many logs in the debug mode.
     debug = False,
-    shard_count = 50,
+    shard_count = most_shards,
     # Takes too long for TSAN. Since this is kind of a stress test that doesn't
     # involve much concurrency, TSAN's usefulness here is limited anyway.
     tags = ["nogotsan"],
@@ -211,7 +212,7 @@ syscall_test(
 
 syscall_test(
     size = "medium",
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:futex_test",
 )
 
@@ -258,7 +259,7 @@ syscall_test(
 
 syscall_test(
     size = "large",
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:itimer_test",
 )
 
@@ -313,7 +314,7 @@ syscall_test(
 
 syscall_test(
     size = "medium",
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:mmap_test",
 )
 
@@ -347,6 +348,7 @@ syscall_test(
 
 syscall_test(
     add_overlay = True,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:open_test",
 )
 
@@ -376,7 +378,7 @@ syscall_test(
 syscall_test(
     size = "large",
     add_overlay = True,
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:pipe_test",
 )
 
@@ -448,7 +450,7 @@ syscall_test(
 
 syscall_test(
     size = "medium",
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:pty_test",
 )
 
@@ -475,6 +477,7 @@ syscall_test(
 )
 
 syscall_test(
+    shard_count = more_shards,
     test = "//test/syscalls/linux:raw_socket_test",
 )
 
@@ -490,7 +493,7 @@ syscall_test(
 
 syscall_test(
     size = "medium",
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:readv_socket_test",
 )
 
@@ -539,7 +542,7 @@ syscall_test(
 )
 
 syscall_test(
-    shard_count = 20,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:semaphore_test",
 )
 
@@ -594,7 +597,7 @@ syscall_test(
 
 syscall_test(
     size = "large",
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_abstract_test",
 )
 
@@ -605,7 +608,7 @@ syscall_test(
 
 syscall_test(
     size = "large",
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_domain_test",
 )
 
@@ -618,19 +621,19 @@ syscall_test(
 syscall_test(
     size = "large",
     add_overlay = True,
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_filesystem_test",
 )
 
 syscall_test(
     size = "large",
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_inet_loopback_test",
 )
 
 syscall_test(
     size = "large",
-    shard_count = 50,
+    shard_count = most_shards,
     # Takes too long for TSAN. Creates a lot of TCP sockets.
     tags = ["nogotsan"],
     test = "//test/syscalls/linux:socket_inet_loopback_nogotsan_test",
@@ -638,7 +641,7 @@ syscall_test(
 
 syscall_test(
     size = "large",
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_ip_tcp_generic_loopback_test",
 )
 
@@ -649,13 +652,13 @@ syscall_test(
 
 syscall_test(
     size = "large",
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_ip_tcp_loopback_test",
 )
 
 syscall_test(
     size = "medium",
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_ip_tcp_udp_generic_loopback_test",
 )
 
@@ -666,7 +669,7 @@ syscall_test(
 
 syscall_test(
     size = "large",
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_ip_udp_loopback_test",
 )
 
@@ -677,6 +680,7 @@ syscall_test(
 
 syscall_test(
     size = "medium",
+    shard_count = more_shards,
     # Takes too long under gotsan to run.
     tags = ["nogotsan"],
     test = "//test/syscalls/linux:socket_ipv4_udp_unbound_loopback_nogotsan_test",
@@ -691,6 +695,7 @@ syscall_test(
 )
 
 syscall_test(
+    shard_count = more_shards,
     test = "//test/syscalls/linux:socket_ip_unbound_test",
 )
 
@@ -753,7 +758,7 @@ syscall_test(
 syscall_test(
     # NOTE(b/116636318): Large sendmsg may stall a long time.
     size = "enormous",
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:socket_unix_dgram_local_test",
 )
 
@@ -765,14 +770,14 @@ syscall_test(
 syscall_test(
     size = "large",
     add_overlay = True,
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_unix_pair_test",
 )
 
 syscall_test(
     # NOTE(b/116636318): Large sendmsg may stall a long time.
     size = "enormous",
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:socket_unix_seqpacket_local_test",
 )
 
@@ -798,13 +803,13 @@ syscall_test(
 
 syscall_test(
     size = "medium",
-    shard_count = 10,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:socket_unix_unbound_seqpacket_test",
 )
 
 syscall_test(
     size = "large",
-    shard_count = 50,
+    shard_count = most_shards,
     test = "//test/syscalls/linux:socket_unix_unbound_stream_test",
 )
 
@@ -858,7 +863,7 @@ syscall_test(
 
 syscall_test(
     size = "medium",
-    shard_count = 10,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:tcp_socket_test",
 )
 
@@ -867,6 +872,7 @@ syscall_test(
 )
 
 syscall_test(
+    shard_count = more_shards,
     test = "//test/syscalls/linux:timerfd_test",
 )
 
@@ -903,7 +909,7 @@ syscall_test(
 syscall_test(
     size = "medium",
     add_hostinet = True,
-    shard_count = 10,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:udp_socket_test",
 )
 
@@ -947,7 +953,7 @@ syscall_test(
 
 syscall_test(
     size = "medium",
-    shard_count = 5,
+    shard_count = more_shards,
     test = "//test/syscalls/linux:wait_test",
 )
 
diff --git a/tools/bazel.mk b/tools/bazel.mk
index 9063eebaf..ca5621a9c 100644
--- a/tools/bazel.mk
+++ b/tools/bazel.mk
@@ -41,6 +41,9 @@ DOCKER_CONFIG := /etc/docker
 # Bazel flags.
 BAZEL := bazel $(STARTUP_OPTIONS)
 BASE_OPTIONS := --color=no --curses=no
+ifneq (,$(BAZEL_CONFIG))
+BASE_OPTIONS += --config=$(BAZEL_CONFIG)
+endif
 
 # Basic options.
 UID := $(shell id -u ${USER})
@@ -103,11 +106,6 @@ FULL_DOCKER_RUN_OPTIONS += --group-add $(KVM_GROUP)
 endif
 endif
 
-# Load the appropriate config.
-ifneq (,$(BAZEL_CONFIG))
-OPTIONS += --config=$(BAZEL_CONFIG)
-endif
-
 bazel-image: load-default
 	@if docker ps --all | grep $(BUILDER_NAME); then docker rm -f $(BUILDER_NAME); fi
 	docker run --user 0:0 --entrypoint "" --name $(BUILDER_NAME) \
diff --git a/tools/bazeldefs/defs.bzl b/tools/bazeldefs/defs.bzl
index c2f94bb9c..51e17a79a 100644
--- a/tools/bazeldefs/defs.bzl
+++ b/tools/bazeldefs/defs.bzl
@@ -7,6 +7,8 @@ build_test = _build_test
 bzl_library = _bzl_library
 rbe_platform = native.platform
 rbe_toolchain = native.toolchain
+more_shards = 4
+most_shards = 8
 
 def short_path(path):
     return path
diff --git a/tools/defs.bzl b/tools/defs.bzl
index 2c8129e7e..b6f188aeb 100644
--- a/tools/defs.bzl
+++ b/tools/defs.bzl
@@ -8,7 +8,7 @@ change for Google-internal and bazel-compatible rules.
 load("//tools/go_stateify:defs.bzl", "go_stateify")
 load("//tools/go_marshal:defs.bzl", "go_marshal", "marshal_deps", "marshal_test_deps")
 load("//tools/nogo:defs.bzl", "nogo_test")
-load("//tools/bazeldefs:defs.bzl", _build_test = "build_test", _bzl_library = "bzl_library", _coreutil = "coreutil", _default_installer = "default_installer", _default_net_util = "default_net_util", _proto_library = "proto_library", _rbe_platform = "rbe_platform", _rbe_toolchain = "rbe_toolchain", _select_arch = "select_arch", _select_system = "select_system", _short_path = "short_path")
+load("//tools/bazeldefs:defs.bzl", _build_test = "build_test", _bzl_library = "bzl_library", _coreutil = "coreutil", _default_installer = "default_installer", _default_net_util = "default_net_util", _more_shards = "more_shards", _most_shards = "most_shards", _proto_library = "proto_library", _rbe_platform = "rbe_platform", _rbe_toolchain = "rbe_toolchain", _select_arch = "select_arch", _select_system = "select_system", _short_path = "short_path")
 load("//tools/bazeldefs:cc.bzl", _cc_binary = "cc_binary", _cc_flags_supplier = "cc_flags_supplier", _cc_grpc_library = "cc_grpc_library", _cc_library = "cc_library", _cc_proto_library = "cc_proto_library", _cc_test = "cc_test", _cc_toolchain = "cc_toolchain", _gbenchmark = "gbenchmark", _grpcpp = "grpcpp", _gtest = "gtest", _vdso_linker_option = "vdso_linker_option")
 load("//tools/bazeldefs:go.bzl", _gazelle = "gazelle", _go_binary = "go_binary", _go_embed_data = "go_embed_data", _go_grpc_and_proto_libraries = "go_grpc_and_proto_libraries", _go_library = "go_library", _go_path = "go_path", _go_proto_library = "go_proto_library", _go_test = "go_test", _select_goarch = "select_goarch", _select_goos = "select_goos")
 load("//tools/bazeldefs:pkg.bzl", _pkg_deb = "pkg_deb", _pkg_tar = "pkg_tar")
@@ -26,6 +26,8 @@ short_path = _short_path
 rbe_platform = _rbe_platform
 rbe_toolchain = _rbe_toolchain
 coreutil = _coreutil
+more_shards = _more_shards
+most_shards = _most_shards
 
 # C++ rules.
 cc_binary = _cc_binary