Set the local CPU reservation for tests based on their "cpu:<n>" tag.

This lets users specify that their test needs a minimum of <n> CPU cores
to run and not be flaky. Example for a reservation of 4 CPUs:

sh_test(
  name = "test",
  size = "large",
  srcs = ["test.sh"],
  tags = ["cpu:4"],
)

This could also be used by remote execution strategies to tune their
resource adjustment.

RELNOTES: You can increase the CPU reservation for tests by adding a "cpu:<n>" (e.g. "cpu:4" for four cores) tag to their rule in a BUILD file. This can be used if tests would otherwise overwhelm your system if there's too much parallelism.
PiperOrigin-RevId: 154856091

6 files changed, 227 insertions, 12 deletions

diff --git a/site/versions/master/docs/test-encyclopedia.html b/site/versions/master/docs/test-encyclopedia.html
index b0eca8bdc2..6a9b0e85e9 100644
--- a/site/versions/master/docs/test-encyclopedia.html
+++ b/site/versions/master/docs/test-encyclopedia.html
@@ -332,7 +332,9 @@ resolve.</p>
 <h4>Other resources</h4>
 
 <p>Tests are granted at least one CPU core. Others may be available but this
-is not guaranteed. Other performance aspects of this core are not specified.</p>
+is not guaranteed. Other performance aspects of this core are not specified.
+You can increase the reservation to a higher number of CPU cores by adding the
+tag "cpu:n" (where n is a positive number) to a test rule.</p>
 
 <p>Tests may create subprocesses, but not process groups or sessions.</p>
 
diff --git a/src/main/java/com/google/devtools/build/lib/BUILD b/src/main/java/com/google/devtools/build/lib/BUILD
index 02086d9fae..7d6d25836b 100644
--- a/src/main/java/com/google/devtools/build/lib/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/BUILD
@@ -516,6 +516,10 @@ java_library(
     visibility = [
         "//src/main/java/com/google/devtools/build/lib/rules:__subpackages__",
     ],
+    deps = [
+        ":util",
+        "//third_party:auto_value",
+    ],
 )
 
 java_library(
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/ExecutionRequirements.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/ExecutionRequirements.java
index 43fa02c843..b929483d75 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/ExecutionRequirements.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/ExecutionRequirements.java
@@ -14,10 +14,126 @@
 
 package com.google.devtools.build.lib.analysis.actions;
 
+import com.google.auto.value.AutoValue;
+import com.google.common.base.Function;
+import com.google.devtools.build.lib.util.Preconditions;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
 /**
  * Strings used to express requirements on action execution environments.
  */
 public class ExecutionRequirements {
+
+  /** An execution requirement that can be split into a key and a value part using a regex. */
+  @AutoValue
+  public abstract static class ParseableRequirement {
+
+    /**
+     * Thrown when a {@link ParseableRequirement} feels responsible for a tag, but the {@link
+     * #validator()} method returns an error.
+     */
+    public static class ValidationException extends Exception {
+      private final String tagValue;
+
+      /**
+       * Creates a new {@link ValidationException}.
+       *
+       * @param tagValue the erroneous value that was parsed from the tag.
+       * @param errorMsg an error message that tells the user what's wrong with the value.
+       */
+      public ValidationException(String tagValue, String errorMsg) {
+        super(errorMsg);
+        this.tagValue = tagValue;
+      }
+
+      /**
+       * Returns the erroneous value of the parsed tag.
+       *
+       * <p>Useful to put in error messages shown to the user.
+       */
+      public String getTagValue() {
+        return tagValue;
+      }
+    }
+
+    /**
+     * Create a new parseable execution requirement definition.
+     *
+     * <p>If a tag doesn't match the detectionPattern, it will be ignored. If a tag matches the
+     * detectionPattern, but not the validationPattern, it is assumed that the value is somehow
+     * wrong (e.g. the user put a float or random string where we expected an integer).
+     *
+     * @param userFriendlyName a human readable name of the tag and its format, e.g. "cpu:<int>"
+     * @param detectionPattern a regex that will be used to detect whether a tag matches this
+     *     execution requirement. It should have one capture group that grabs the value of the tag.
+     *     This should be general enough to permit even wrong value types. Example: "cpu:(.+)".
+     * @param validator a Function that will be used to validate the value of the tag. It should
+     *     return null if the value is fine to use or a human-friendly error message describing why
+     *     the value is not valid.
+     */
+    static ParseableRequirement create(
+        String userFriendlyName, Pattern detectionPattern, Function<String, String> validator) {
+      return new AutoValue_ExecutionRequirements_ParseableRequirement(
+          userFriendlyName, detectionPattern, validator);
+    }
+
+    public abstract String userFriendlyName();
+
+    public abstract Pattern detectionPattern();
+
+    public abstract Function<String, String> validator();
+
+    /**
+     * Returns the parsed value from a tag, if this {@link ParseableRequirement} detects that it is
+     * responsible for it, otherwise returns {@code null}.
+     *
+     * @throws ValidationException if the value parsed out of the tag doesn't pass the validator.
+     */
+    public String parseIfMatches(String tag) throws ValidationException {
+      Matcher matcher = detectionPattern().matcher(tag);
+      if (!matcher.matches()) {
+        return null;
+      }
+      String tagValue = matcher.group(1);
+      String errorMsg = validator().apply(tagValue);
+      if (errorMsg != null) {
+        throw new ValidationException(tagValue, errorMsg);
+      }
+      return tagValue;
+    }
+  }
+
   /** If an action would not successfully run other than on Darwin. */
   public static final String REQUIRES_DARWIN = "requires-darwin";
+
+  /** How many hardware threads an action requires for execution. */
+  public static final ParseableRequirement CPU =
+      ParseableRequirement.create(
+          "cpu:<int>",
+          Pattern.compile("cpu:(.+)"),
+          new Function<String, String>() {
+            @Override
+            public String apply(String s) {
+              Preconditions.checkNotNull(s);
+
+              int value;
+              try {
+                value = Integer.parseInt(s);
+              } catch (NumberFormatException e) {
+                return "can't be parsed as an integer";
+              }
+
+              // De-and-reserialize & compare to only allow canonical integer formats.
+              if (!Integer.toString(value).equals(s)) {
+                return "must be in canonical format (e.g. '4' instead of '+04')";
+              }
+
+              if (value < 1) {
+                return "can't be zero or negative";
+              }
+
+              return null;
+            }
+          });
 }
diff --git a/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java b/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java
index f63d5196fb..58aed286d0 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java
@@ -22,6 +22,7 @@ import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ExecutionStrategy;
 import com.google.devtools.build.lib.actions.Executor;
+import com.google.devtools.build.lib.actions.ResourceSet;
 import com.google.devtools.build.lib.actions.SimpleSpawn;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnActionContext;
@@ -114,6 +115,12 @@ public class StandaloneTestStrategy extends TestStrategy {
     info.put("timeout", "" + getTimeout(action));
     info.putAll(action.getTestProperties().getExecutionInfo());
 
+    ResourceSet localResourceUsage =
+        action
+            .getTestProperties()
+            .getLocalResourceUsage(
+                action.getOwner().getLabel(), executionOptions.usingLocalTestJobs());
+
     Spawn spawn =
         new SimpleSpawn(
             action,
@@ -122,13 +129,11 @@ public class StandaloneTestStrategy extends TestStrategy {
             ImmutableMap.copyOf(info),
             new RunfilesSupplierImpl(
                 runfilesDir.relativeTo(execRoot), action.getExecutionSettings().getRunfiles()),
-            /*inputs=*/ImmutableList.copyOf(action.getInputs()),
-            /*tools=*/ImmutableList.<Artifact>of(),
-            /*filesetManifests=*/ImmutableList.<Artifact>of(),
+            /*inputs=*/ ImmutableList.copyOf(action.getInputs()),
+            /*tools=*/ ImmutableList.<Artifact>of(),
+            /*filesetManifests=*/ ImmutableList.<Artifact>of(),
             ImmutableList.copyOf(action.getSpawnOutputs()),
-            action
-                .getTestProperties()
-                .getLocalResourceUsage(executionOptions.usingLocalTestJobs()));
+            localResourceUsage);
 
     Executor executor = actionExecutionContext.getExecutor();
 
diff --git a/src/main/java/com/google/devtools/build/lib/rules/test/TestTargetProperties.java b/src/main/java/com/google/devtools/build/lib/rules/test/TestTargetProperties.java
index 839a6fee73..ce493d2f0d 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/test/TestTargetProperties.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/test/TestTargetProperties.java
@@ -18,14 +18,17 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.actions.ResourceSet;
 import com.google.devtools.build.lib.actions.Spawn;
+import com.google.devtools.build.lib.actions.UserExecException;
 import com.google.devtools.build.lib.analysis.RuleContext;
+import com.google.devtools.build.lib.analysis.actions.ExecutionRequirements;
+import com.google.devtools.build.lib.analysis.actions.ExecutionRequirements.ParseableRequirement.ValidationException;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.packages.TestSize;
 import com.google.devtools.build.lib.packages.TestTimeout;
 import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.util.Preconditions;
-
 import java.util.List;
 import java.util.Map;
 
@@ -130,10 +133,45 @@ public class TestTargetProperties {
     return isExternal;
   }
 
-  public ResourceSet getLocalResourceUsage(boolean usingLocalTestJobs) {
-    return usingLocalTestJobs
-        ? LOCAL_TEST_JOBS_BASED_RESOURCES
-        : TestTargetProperties.getResourceSetFromSize(size);
+  public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs)
+      throws UserExecException {
+    if (usingLocalTestJobs) {
+      return LOCAL_TEST_JOBS_BASED_RESOURCES;
+    }
+
+    ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size);
+
+    // Tests can override their CPU reservation with a "cpus:<n>" tag.
+    ResourceSet testResourcesFromTag = null;
+    for (String tag : executionInfo.keySet()) {
+      try {
+        String cpus = ExecutionRequirements.CPU.parseIfMatches(tag);
+        if (cpus != null) {
+          if (testResourcesFromTag != null) {
+            throw new UserExecException(
+                String.format(
+                    "%s has more than one '%s' tag, but duplicate tags aren't allowed",
+                    label, ExecutionRequirements.CPU.userFriendlyName()));
+          }
+          testResourcesFromTag =
+              ResourceSet.create(
+                  testResourcesFromSize.getMemoryMb(),
+                  Float.parseFloat(cpus),
+                  testResourcesFromSize.getIoUsage(),
+                  testResourcesFromSize.getLocalTestCount());
+        }
+      } catch (ValidationException e) {
+        throw new UserExecException(
+            String.format(
+                "%s has a '%s' tag, but its value '%s' didn't pass validation: %s",
+                label,
+                ExecutionRequirements.CPU.userFriendlyName(),
+                e.getTagValue(),
+                e.getMessage()));
+      }
+    }
+
+    return testResourcesFromTag != null ? testResourcesFromTag : testResourcesFromSize;
   }
 
   /**
diff --git a/src/test/java/com/google/devtools/build/lib/rules/test/TestTargetPropertiesTest.java b/src/test/java/com/google/devtools/build/lib/rules/test/TestTargetPropertiesTest.java
new file mode 100644
index 0000000000..2a12f56d0b
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/rules/test/TestTargetPropertiesTest.java
@@ -0,0 +1,50 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.rules.test;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.devtools.build.lib.actions.ResourceSet;
+import com.google.devtools.build.lib.analysis.ConfiguredTarget;
+import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link TestTargetProperties}. */
+@RunWith(JUnit4.class)
+public class TestTargetPropertiesTest extends BuildViewTestCase {
+  @Test
+  public void testTestWithCpusTagHasCorrectLocalResourcesEstimate() throws Exception {
+    scratch.file("tests/test.sh", "#!/bin/bash", "exit 0");
+    scratch.file(
+        "tests/BUILD",
+        "sh_test(",
+        "  name = 'test',",
+        "  size = 'small',",
+        "  srcs = ['test.sh'],",
+        "  tags = ['cpu:4'],",
+        ")");
+    ConfiguredTarget testTarget = getConfiguredTarget("//tests:test");
+    TestRunnerAction testAction =
+        (TestRunnerAction)
+            getGeneratingAction(TestProvider.getTestStatusArtifacts(testTarget).get(0));
+    ResourceSet localResourceUsage =
+        testAction
+            .getTestProperties()
+            .getLocalResourceUsage(testAction.getOwner().getLabel(), false);
+    assertThat(localResourceUsage.getCpuUsage()).isEqualTo(4.0);
+  }
+}