Only cache runfiles mappings of tests if they have more than one shard or run per test.

This is an attempt to fix Bazel excessively OOMing at Google. This change was behavior was caused by https://github.com/bazelbuild/bazel/commit/08df6f1ebeaccb43d3aac1058f8d9e28e11acd2b where I believed that the soft reference makes the caching a no-op wrt. memory use.

This does not quite reinstate the same behavior as before https://github.com/bazelbuild/bazel/commit/08df6f1ebeaccb43d3aac1058f8d9e28e11acd2b: firstly, it doesn't take --test_sharding_strategy into account (it could, but I would have had to make TestShardingStrategy public instead of package-internal) and it doesn't garbage collect the soft reference after the all actions of a given test target finish running. It's still an improvement, though, and profiling shows that it has a good chance of fixing the problem.

RELNOTES: None.
PiperOrigin-RevId: 602689571
Change-Id: Ib68c1b88bf47898ae8a59866447d7790bb6eceee
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/RunfilesSupport.java b/src/main/java/com/google/devtools/build/lib/analysis/RunfilesSupport.java
index 70aafc2..bfc2883 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/RunfilesSupport.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/RunfilesSupport.java
@@ -29,6 +29,7 @@
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue.RunfileSymlinksMode;
 import com.google.devtools.build.lib.analysis.config.RunUnder;
+import com.google.devtools.build.lib.analysis.test.TestConfiguration;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.packages.TargetUtils;
@@ -177,6 +178,35 @@
   private final CommandLine args;
   private final ActionEnvironment actionEnvironment;
 
+  private static boolean cacheRunfilesMappings(RuleContext ruleContext) {
+    // The algorithm: only cache runfiles if there is more than one test runner action.
+    // This is because soft references interact with GC in an unexpected way and it's sometimes even
+    // possible to OOM before soft references are collected. Therefore, let's not pay for their
+    // overhead if they are useless (which they are if there is only one test action)
+    if (!TargetUtils.isTestRule(ruleContext.getTarget())) {
+      return false;
+    }
+
+    TestConfiguration testConfiguration = ruleContext.getFragment(TestConfiguration.class);
+    if (testConfiguration == null) {
+      return false;
+    }
+
+    int runsPerTest = testConfiguration.getRunsPerTestForLabel(ruleContext.getLabel());
+    if (runsPerTest > 1) {
+      return true;
+    }
+
+    if (!ruleContext.attributes().has("shard_count", Type.INTEGER)) {
+      return false;
+    }
+
+    int explicitShardCount =
+        ruleContext.attributes().get("shard_count", Type.INTEGER).toIntUnchecked();
+
+    return explicitShardCount > 1;
+  }
+
   /**
    * Creates the RunfilesSupport helper with the given executable and runfiles.
    *
@@ -195,7 +225,6 @@
         ruleContext.getConfiguration().getRunfileSymlinksMode();
     boolean buildRunfileManifests = ruleContext.getConfiguration().buildRunfileManifests();
     boolean buildRunfileLinks = ruleContext.getConfiguration().buildRunfileLinks();
-    boolean cacheRunfilesMapping = TargetUtils.isTestRule(ruleContext.getTarget());
 
     // Adding run_under target to the runfiles manifest so it would become part
     // of runfiles tree and would be executable everywhere.
@@ -239,7 +268,7 @@
             runfiles,
             repoMappingManifest,
             buildRunfileLinks,
-            cacheRunfilesMapping,
+            cacheRunfilesMappings(ruleContext),
             runfileSymlinksMode);
 
     Artifact runfilesMiddleman =