Implement extra resources for test actions

Adds support for specifying extra resources available for the bazel
invocation and requirement for test actions. Using
--local_extra_resources=<resource name>=<amount of resource> one can
specify the resources available for tests. Specifying a tag
"resources:<resource name>:<amount of resources>" for a test action will
tell bazel how much of the said resource this test will require. Bazel
will then limit concurrently running test actions using the same
resource if there isn't enough of said resource available.

Author: exoson

src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java

@@ -152,6 +152,35 @@ public String parseIfMatches(String tag) throws ValidationException {
             return null;
           });
 
+  /** How many extra resources an action requires for execution. */
+  public static final ParseableRequirement RESOURCES =
+      ParseableRequirement.create(
+          "resources:<str>:<int>",
+          Pattern.compile("resources:(.+:.+)"),
+          s -> {
+            Preconditions.checkNotNull(s);
+
+            int splitIndex = s.indexOf(":");
+            String resourceCount = s.substring(splitIndex+1);
+            int value;
+            try {
+              value = Integer.parseInt(resourceCount);
+            } 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(resourceCount)) {
+              return "must be in canonical format (e.g. '4' instead of '+04')";
+            }
+
+            if (value < 1) {
+              return "can't be zero or negative";
+            }
+
+            return null;
+          });
+
   /** If an action supports running in persistent worker mode. */
   public static final String SUPPORTS_WORKERS = "supports-workers";
 

src/main/java/com/google/devtools/build/lib/actions/LocalHostResourceFallback.java

@@ -27,6 +27,7 @@ public class LocalHostResourceFallback {
       ResourceSet.create(
           3.0 * (Runtime.getRuntime().maxMemory() >> 20),
           Runtime.getRuntime().availableProcessors(),
+          null,
           Integer.MAX_VALUE);
 
   public static ResourceSet getLocalHostResources() {

src/main/java/com/google/devtools/build/lib/actions/LocalHostResourceManagerDarwin.java

@@ -44,6 +44,7 @@ public static ResourceSet getLocalHostResources() {
       return ResourceSet.create(
           ramMb,
           logicalCpuCount,
+          null,
           Integer.MAX_VALUE);
     } catch (IOException e) {
       return null;

src/main/java/com/google/devtools/build/lib/actions/LocalHostResourceManagerLinux.java

@@ -48,6 +48,7 @@ public static ResourceSet getLocalHostResources() {
       return ResourceSet.create(
           ramMb,
           logicalCpuCount,
+          null,
           Integer.MAX_VALUE);
     } catch (IOException e) {
       return null;

src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java

@@ -25,9 +25,13 @@
 import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.util.Pair;
 import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 
 /**
@@ -137,6 +141,10 @@ public static ResourceManager instance() {
   // definition in the ResourceSet class.
   private double usedRam;
 
+  // Used amount of extra resources. Corresponds to the extra resource
+  // definition in the ResourceSet class.
+  private Map<String, Float> usedExtraResources;
+
   // Used local test count. Corresponds to the local test count definition in the ResourceSet class.
   private int usedLocalTestCount;
 
@@ -145,6 +153,7 @@ public static ResourceManager instance() {
 
   private ResourceManager() {
     requestList = new LinkedList<>();
+    usedExtraResources = new HashMap<>();
   }
 
   @VisibleForTesting public static ResourceManager instanceForTestingOnly() {
@@ -158,6 +167,7 @@ private ResourceManager() {
   public synchronized void resetResourceUsage() {
     usedCpu = 0;
     usedRam = 0;
+    usedExtraResources = new HashMap<>();
     usedLocalTestCount = 0;
     for (Pair<ResourceSet, CountDownLatch> request : requestList) {
       // CountDownLatch can be set only to 0 or 1.
@@ -177,6 +187,7 @@ public synchronized void setAvailableResources(ResourceSet resources) {
         ResourceSet.create(
             staticResources.getMemoryMb(),
             staticResources.getCpuUsage(),
+            staticResources.getExtraResourceUsage(),
             staticResources.getLocalTestCount());
     processWaitingThreads();
   }
@@ -264,14 +275,24 @@ ResourceHandle tryAcquire(ActionExecutionMetadata owner, ResourceSet resources)
   private void incrementResources(ResourceSet resources) {
     usedCpu += resources.getCpuUsage();
     usedRam += resources.getMemoryMb();
+
+    for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
+      String key = (String)resource.getKey();
+      float value = resource.getValue();
+      if (usedExtraResources.containsKey(key)) {
+        value += (float)usedExtraResources.get(key);
+      }
+      usedExtraResources.put(key, value);
+    }
+
     usedLocalTestCount += resources.getLocalTestCount();
   }
 
   /**
    * Return true if any resources have been claimed through this manager.
    */
   public synchronized boolean inUse() {
-    return usedCpu != 0.0 || usedRam != 0.0 || usedLocalTestCount != 0 || !requestList.isEmpty();
+    return usedCpu != 0.0 || usedRam != 0.0 || !usedExtraResources.isEmpty() || usedLocalTestCount != 0 || !requestList.isEmpty();
   }
 
 
@@ -322,6 +343,13 @@ private synchronized CountDownLatch acquire(ResourceSet resources) {
   private synchronized boolean release(ResourceSet resources) {
     usedCpu -= resources.getCpuUsage();
     usedRam -= resources.getMemoryMb();
+
+    for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
+      String key = (String)resource.getKey();
+      float value = (float)usedExtraResources.get(key) - resource.getValue();
+      usedExtraResources.put(key, value);
+    }
+
     usedLocalTestCount -= resources.getLocalTestCount();
 
     // TODO(bazel-team): (2010) rounding error can accumulate and value below can end up being
@@ -333,6 +361,19 @@ private synchronized boolean release(ResourceSet resources) {
     if (usedRam < epsilon) {
       usedRam = 0;
     }
+
+    Set<String> toRemove = new HashSet<>();
+    for (Map.Entry<String, Float> resource : usedExtraResources.entrySet()) {
+      String key = (String)resource.getKey();
+      float value = (float)usedExtraResources.get(key);
+      if (value < epsilon) {
+        toRemove.add(key);
+      }
+    }
+    for (String key : toRemove) {
+      usedExtraResources.remove(key);
+    }
+
     if (!requestList.isEmpty()) {
       processWaitingThreads();
       return true;
@@ -365,7 +406,7 @@ private boolean areResourcesAvailable(ResourceSet resources) {
     Preconditions.checkNotNull(availableResources);
     // Comparison below is robust, since any calculation errors will be fixed
     // by the release() method.
-    if (usedCpu == 0.0 && usedRam == 0.0 && usedLocalTestCount == 0) {
+    if (usedCpu == 0.0 && usedRam == 0.0 && usedExtraResources.isEmpty() && usedLocalTestCount == 0) {
       return true;
     }
     // Use only MIN_NECESSARY_???_RATIO of the resource value to check for
@@ -409,9 +450,19 @@ private boolean areResourcesAvailable(ResourceSet resources) {
     // 3) If used resource amount is less than total available resource amount.
     boolean cpuIsAvailable = cpu == 0.0 || usedCpu == 0.0 || usedCpu + cpu <= availableCpu;
     boolean ramIsAvailable = ram == 0.0 || usedRam == 0.0 || ram <= remainingRam;
+    boolean extraResourcesIsAvailable = true;
+    for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
+      String key = (String)resource.getKey();
+      float used = (float)usedExtraResources.getOrDefault(key, 0f);
+      float requested = resource.getValue();
+      float available = (float)availableResources.getExtraResourceUsage().getOrDefault(key, 0f);
+      if (requested != 0.0 && used != 0.0 && requested + used > available) {
+        extraResourcesIsAvailable = false;
+      }
+    }
     boolean localTestCountIsAvailable = localTestCount == 0 || usedLocalTestCount == 0
         || usedLocalTestCount + localTestCount <= availableLocalTestCount;
-    return cpuIsAvailable && ramIsAvailable && localTestCountIsAvailable;
+    return cpuIsAvailable && ramIsAvailable && extraResourcesIsAvailable && localTestCountIsAvailable;
   }
 
   @VisibleForTesting
@@ -421,6 +472,6 @@ synchronized int getWaitCount() {
 
   @VisibleForTesting
   synchronized boolean isAvailable(double ram, double cpu, int localTestCount) {
-    return areResourcesAvailable(ResourceSet.create(ram, cpu, localTestCount));
+    return areResourcesAvailable(ResourceSet.create(ram, cpu, null, localTestCount));
   }
 }

src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java

@@ -19,7 +19,10 @@
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.common.options.Converter;
 import com.google.devtools.common.options.OptionsParsingException;
+import java.util.HashMap;
 import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
 import java.util.NoSuchElementException;
 
 /**
@@ -41,12 +44,24 @@ public class ResourceSet {
   /** The number of CPUs, or fractions thereof. */
   private final double cpuUsage;
 
+  /** Map of extra resources mapping name of the resource to a value. */
+  private final Map<String, Float> extraResourceUsage;
+
   /** The number of local tests. */
   private final int localTestCount;
 
   private ResourceSet(double memoryMb, double cpuUsage, int localTestCount) {
+    this(memoryMb, cpuUsage, null, localTestCount);
+  }
+
+  private ResourceSet(double memoryMb, double cpuUsage, Map<String, Float> extraResourceUsage, int localTestCount) {
     this.memoryMb = memoryMb;
     this.cpuUsage = cpuUsage;
+    if (extraResourceUsage == null) {
+        this.extraResourceUsage = new HashMap<>();
+    } else {
+        this.extraResourceUsage = extraResourceUsage;
+    }
     this.localTestCount = localTestCount;
   }
 
@@ -79,11 +94,11 @@ public static ResourceSet createWithLocalTestCount(int localTestCount) {
    */
   @AutoCodec.Instantiator
   public static ResourceSet create(
-      double memoryMb, double cpuUsage, int localTestCount) {
-    if (memoryMb == 0 && cpuUsage == 0 && localTestCount == 0) {
+      double memoryMb, double cpuUsage, Map<String, Float> extraResourceUsage, int localTestCount) {
+    if (memoryMb == 0 && cpuUsage == 0 && (extraResourceUsage == null || extraResourceUsage.size() == 0) && localTestCount == 0) {
       return ZERO;
     }
-    return new ResourceSet(memoryMb, cpuUsage, localTestCount);
+    return new ResourceSet(memoryMb, cpuUsage, extraResourceUsage, localTestCount);
   }
 
   /** Returns the amount of real memory (resident set size) used in MB. */
@@ -103,16 +118,25 @@ public double getCpuUsage() {
     return cpuUsage;
   }
 
+  public Map<String, Float> getExtraResourceUsage() {
+    return extraResourceUsage;
+  }
+
   /** Returns the local test count used. */
   public int getLocalTestCount() {
     return localTestCount;
   }
 
   @Override
   public String toString() {
+    StringBuilder sb = new StringBuilder();
+    for (Map.Entry<String, Float> resource : extraResourceUsage.entrySet()) {
+      sb.append(resource.getKey() + ": " + resource.getValue() + "\n");
+    }
     return "Resources: \n"
         + "Memory: " + memoryMb + "M\n"
         + "CPU: " + cpuUsage + "\n"
+        + sb.toString()
         + "Local tests: " + localTestCount + "\n";
   }
 
@@ -135,7 +159,7 @@ public ResourceSet convert(String input) throws OptionsParsingException {
         if (memoryMb <= 0.0 || cpuUsage <= 0.0) {
           throw new OptionsParsingException("All resource values must be positive");
         }
-        return create(memoryMb, cpuUsage, Integer.MAX_VALUE);
+        return create(memoryMb, cpuUsage, null, Integer.MAX_VALUE);
       } catch (NumberFormatException | NoSuchElementException nfe) {
         throw new OptionsParsingException("Expected exactly 3 comma-separated float values", nfe);
       }

src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java

@@ -31,6 +31,7 @@
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
 import com.google.devtools.build.lib.server.FailureDetails.TestAction;
 import com.google.devtools.build.lib.server.FailureDetails.TestAction.Code;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -46,10 +47,10 @@ public class TestTargetProperties {
    * <p>When changing these values, remember to update the documentation at
    * attributes/test/size.html.
    */
-  private static final ResourceSet SMALL_RESOURCES = ResourceSet.create(20, 1, 1);
-  private static final ResourceSet MEDIUM_RESOURCES = ResourceSet.create(100, 1, 1);
-  private static final ResourceSet LARGE_RESOURCES = ResourceSet.create(300, 1, 1);
-  private static final ResourceSet ENORMOUS_RESOURCES = ResourceSet.create(800, 1, 1);
+  private static final ResourceSet SMALL_RESOURCES = ResourceSet.create(20, 1, null, 1);
+  private static final ResourceSet MEDIUM_RESOURCES = ResourceSet.create(100, 1, null, 1);
+  private static final ResourceSet LARGE_RESOURCES = ResourceSet.create(300, 1, null, 1);
+  private static final ResourceSet ENORMOUS_RESOURCES = ResourceSet.create(800, 1, null, 1);
   private static final ResourceSet LOCAL_TEST_JOBS_BASED_RESOURCES =
       ResourceSet.createWithLocalTestCount(1);
 
@@ -150,23 +151,21 @@ public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs
     ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size);
 
     // Tests can override their CPU reservation with a "cpus:<n>" tag.
-    ResourceSet testResourcesFromTag = null;
+    // Tests can also specify requirements for extra resources using "resources:<resource name>:<n>" tag.
+    double cpuCount = -1.0;
+    Map<String, Float> extraResources = new HashMap<>();
     for (String tag : executionInfo.keySet()) {
       try {
         String cpus = ExecutionRequirements.CPU.parseIfMatches(tag);
         if (cpus != null) {
-          if (testResourcesFromTag != null) {
+          if (cpuCount != -1.0) {
             String message =
                 String.format(
                     "%s has more than one '%s' tag, but duplicate tags aren't allowed",
                     label, ExecutionRequirements.CPU.userFriendlyName());
             throw new UserExecException(createFailureDetail(message, Code.DUPLICATE_CPU_TAGS));
           }
-          testResourcesFromTag =
-              ResourceSet.create(
-                  testResourcesFromSize.getMemoryMb(),
-                  Float.parseFloat(cpus),
-                  testResourcesFromSize.getLocalTestCount());
+          cpuCount = Float.parseFloat(cpus);
         }
       } catch (ValidationException e) {
         String message =
@@ -178,9 +177,37 @@ public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs
                 e.getMessage());
         throw new UserExecException(createFailureDetail(message, Code.INVALID_CPU_TAG));
       }
+      try {
+        String extras = ExecutionRequirements.RESOURCES.parseIfMatches(tag);
+        if (extras != null) {
+          int splitIndex = extras.indexOf(":");
+          String resourceName = extras.substring(0, splitIndex);
+          String resourceCount = extras.substring(splitIndex+1);
+          if (extraResources.get(resourceName) != null) {
+            String message =
+                String.format(
+                    "%s has more than one '%s' tag, but duplicate tags aren't allowed",
+                    label, ExecutionRequirements.RESOURCES.userFriendlyName());
+            throw new UserExecException(createFailureDetail(message, Code.DUPLICATE_CPU_TAGS));
+          }
+          extraResources.put(resourceName, Float.parseFloat(resourceCount));
+        }
+      } catch (ValidationException e) {
+        String message =
+            String.format(
+                "%s has a '%s' tag, but its value '%s' didn't pass validation: %s",
+                label,
+                ExecutionRequirements.RESOURCES.userFriendlyName(),
+                e.getTagValue(),
+                e.getMessage());
+        throw new UserExecException(createFailureDetail(message, Code.INVALID_CPU_TAG));
+      }
     }
-
-    return testResourcesFromTag != null ? testResourcesFromTag : testResourcesFromSize;
+    return ResourceSet.create(
+            testResourcesFromSize.getMemoryMb(),
+            cpuCount != -1.0 ? cpuCount : testResourcesFromSize.getCpuUsage(),
+            extraResources,
+            testResourcesFromSize.getLocalTestCount());
   }
 
   private static FailureDetail createFailureDetail(String message, Code detailedCode) {