fixup Implement extra resources for test actions

Addressed comments on bazelbuild#13996
Fixed issues in tests and built and tested with
lowRISC/opentitan#16436

Signed-off-by: Drew Macrae <[email protected]>

Author: Drew Macrae

src/main/java/com/google/devtools/build/lib/actions/BUILD

@@ -134,7 +134,6 @@ java_library(
         "//third_party:flogger",
         "//third_party:guava",
         "//third_party:jsr305",
-        "//third_party:rxjava3",
         "//third_party/protobuf:protobuf_java",
     ],
 )
@@ -319,7 +318,6 @@ java_library(
         "//third_party:flogger",
         "//third_party:guava",
         "//third_party:jsr305",
-        "//third_party:rxjava3",
     ],
 )
 

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

@@ -157,27 +157,22 @@ public String parseIfMatches(String tag) throws ValidationException {
   /** How many extra resources an action requires for execution. */
   public static final ParseableRequirement RESOURCES =
       ParseableRequirement.create(
-          "resources:<str>:<int>",
+          "resources:<str>:<float>",
           Pattern.compile("resources:(.+:.+)"),
           s -> {
             Preconditions.checkNotNull(s);
 
             int splitIndex = s.indexOf(":");
             String resourceCount = s.substring(splitIndex+1);
-            int value;
+            float value;
             try {
-              value = Integer.parseInt(resourceCount);
+              value = Float.parseFloat(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')";
+              return "can't be parsed as a float";
             }
 
-            if (value < 1) {
-              return "can't be zero or negative";
+            if (value < 0) {
+              return "can't be negative";
             }
 
             return null;

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

@@ -32,6 +32,7 @@
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import javax.annotation.Nullable;
@@ -197,10 +198,6 @@ public double getUsedCPU() {
   /** If set, local-only actions are given priority over dynamically run actions. */
   private boolean prioritizeLocalActions;
 
-  private ResourceManager() {
-    usedExtraResources = new HashMap<>();
-  }
-
   @VisibleForTesting
   public static ResourceManager instanceForTestingOnly() {
     return new ResourceManager();
@@ -392,7 +389,7 @@ public void acquireResourceOwnership() {
    * wait.
    */
   private synchronized LatchWithWorker acquire(ResourceSet resources, ResourcePriority priority)
-      throws IOException, InterruptedException {
+      throws IOException, InterruptedException, NoSuchElementException {
     if (areResourcesAvailable(resources)) {
       Worker worker = incrementResources(resources);
       return new LatchWithWorker(/* latch= */ null, worker);
@@ -441,12 +438,6 @@ private synchronized void releaseResourcesOnly(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
@@ -460,14 +451,14 @@ private synchronized void releaseResourcesOnly(ResourceSet resources) {
     }
 
     Set<String> toRemove = new HashSet<>();
-    usedExtraResources.entrySet().forEach(
-      resource -> {
-        String key = (String)resource.getKey();
-        float value = (float)usedExtraResources.get(key);
-        if (value < epsilon) {
-          toRemove.add(key);
-        }
-      });
+    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);
+      if (value < epsilon) {
+        toRemove.add(key);
+      }
+    }
     for (String key : toRemove) {
       usedExtraResources.remove(key);
     }
@@ -509,15 +500,28 @@ private synchronized void processWaitingThreads(Deque<Pair<ResourceSet, LatchWit
     }
   }
 
+  /**
+   * Throws an exception if requested extra resource isn't being tracked
+   */
+  private void assertExtraResourcesTracked(ResourceSet resources)
+      throws NoSuchElementException {
+    for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
+      String key = (String)resource.getKey();
+      if (!availableResources.getExtraResourceUsage().containsKey(key)) {
+        throw new NoSuchElementException("Resource "+key+" is not tracked in this resource set.");
+      }
+    }
+  }
+
   /**
    * Return true iff all requested extra resources are considered to be available.
    */
-  private boolean areExtraResourcesAvailable(ResourceSet resources) {
+  private boolean areExtraResourcesAvailable(ResourceSet resources) throws NoSuchElementException {
     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);
+      float available = availableResources.getExtraResourceUsage().get(key);
       float epsilon = 0.0001f;  // Account for possible rounding errors.
       if (requested != 0.0 && used != 0.0 && requested + used > available + epsilon) {
         return false;
@@ -528,7 +532,7 @@ private boolean areExtraResourcesAvailable(ResourceSet resources) {
 
   // Method will return true if all requested resources are considered to be available.
   @VisibleForTesting
-  boolean areResourcesAvailable(ResourceSet resources) {
+  boolean areResourcesAvailable(ResourceSet resources) throws NoSuchElementException {
     Preconditions.checkNotNull(availableResources);
     // Comparison below is robust, since any calculation errors will be fixed
     // by the release() method.
@@ -544,6 +548,10 @@ boolean areResourcesAvailable(ResourceSet resources) {
         workerKey == null
             || (activeWorkers < availableWorkers && workerPool.couldBeBorrowed(workerKey));
 
+    // We test for tracking of extra resources whenever acquired and throw an
+    // exception before acquiring any untracked resource.
+    assertExtraResourcesTracked(resources);
+
     if (usedCpu == 0.0 && usedRam == 0.0 && usedExtraResources.isEmpty() && usedLocalTestCount == 0 && workerIsAvailable) {
       return true;
     }

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

@@ -23,7 +23,6 @@
 import com.google.devtools.build.lib.worker.WorkerKey;
 import com.google.devtools.common.options.Converter;
 import com.google.devtools.common.options.OptionsParsingException;
-import io.reactivex.rxjava3.annotations.NonNull;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 import javax.annotation.Nullable;
@@ -58,11 +57,12 @@ public class ResourceSet implements ResourceSetOrBuilder {
   /** The workerKey of used worker. Null if no worker is used. */
   @Nullable private final WorkerKey workerKey;
 
-  private ResourceSet(double memoryMb, double cpuUsage, int localTestCount, @Nullable WorkerKey workerKey) {
+  private ResourceSet(
+      double memoryMb, double cpuUsage, int localTestCount, @Nullable WorkerKey workerKey) {
     this(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, workerKey);
   }
 
-  private ResourceSet(double memoryMb, double cpuUsage, @NonNull ImmutableMap<String, Float> extraResourceUsage, int localTestCount, @Nullable WorkerKey workerKey) {
+  private ResourceSet(double memoryMb, double cpuUsage, @Nullable ImmutableMap<String, Float> extraResourceUsage, int localTestCount, @Nullable WorkerKey workerKey) {
     this.memoryMb = memoryMb;
     this.cpuUsage = cpuUsage;
     this.extraResourceUsage = extraResourceUsage;
@@ -102,7 +102,7 @@ public static ResourceSet createWithLocalTestCount(int localTestCount) {
    * represent available resources.
    */
   public static ResourceSet create(double memoryMb, double cpuUsage, int localTestCount) {
-    return ResourceSet.create(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, /* wolkerKey= */ null);
+    return ResourceSet.createWithWorkerKey(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, /* wolkerKey= */ null);
   }
 
   /**
@@ -112,7 +112,11 @@ public static ResourceSet create(double memoryMb, double cpuUsage, int localTest
    * ResourceSets that represent available resources.
    */
   public static ResourceSet create(double memoryMb, double cpuUsage, ImmutableMap<String, Float> extraResourceUsage, int localTestCount) {
-    return createWithWorkerKey(memoryMb, cpuUsage, extraResourceUseage, localTestCount, /* workerKey= */ null);
+    return createWithWorkerKey(memoryMb, cpuUsage, extraResourceUsage, localTestCount, /* workerKey= */ null);
+  }
+
+  public static ResourceSet createWithWorkerKey(double memoryMb, double cpuUsage, int localTestCount, WorkerKey workerKey) {
+    return ResourceSet.createWithWorkerKey(memoryMb, cpuUsage, /* extraResourceUsage= */ImmutableMap.of(), localTestCount, workerKey);
   }
 
   public static ResourceSet createWithWorkerKey(

src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java

@@ -929,7 +929,7 @@ public static void configureResourceManager(ResourceManager resourceMgr, BuildRe
     ExecutionOptions options = request.getOptions(ExecutionOptions.class);
     resourceMgr.setPrioritizeLocalActions(options.prioritizeLocalActions);
     ImmutableMap<String, Float> extraResources = options.localExtraResources.stream().collect(
-      ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (v1,v2) -> v1));
+      ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (v1,v2) -> v2));
 
     resourceMgr.setAvailableResources(
         ResourceSet.create(

src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java

@@ -334,11 +334,11 @@ public boolean shouldMaterializeParamFiles() {
       help =
           "Set the number of extra resources available to Bazel. "
             + "Takes in a string-float pair. Can be used multiple times to specify multiple "
-            + "types of extra resources. Bazel will limit concurrently running test actions "
-            + "based on the available extra resources and the extra resources required "
-            + "by the test actions.  Tests can declare the amount of extra resources they need "
+            + "types of extra resources. Bazel will limit concurrently running actions "
+            + "based on the available extra resources and the extra resources required. "
+            + "Tests can declare the amount of extra resources they need "
             + "by using a tag of the \"resources:<resoucename>:<amount>\" format. "
-            + "Available CPU, RAM and test job resources cannot be set with this flag.",
+            + "Available CPU, RAM and resources cannot be set with this flag.",
       converter = Converters.StringToFloatAssignmentConverter.class)
   public List<Map.Entry<String, Float>> localExtraResources;
 

src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnRunner.java

@@ -412,6 +412,7 @@ WorkResponse execInWorkerWorkerAsResource(
         ResourceSet.createWithWorkerKey(
             spawn.getLocalResources().getMemoryMb(),
             spawn.getLocalResources().getCpuUsage(),
+            spawn.getLocalResources().getExtraResourceUsage(),
             spawn.getLocalResources().getLocalTestCount(),
             key);
 

src/main/java/com/google/devtools/common/options/Converters.java

@@ -475,7 +475,7 @@ public String getTypeDescription() {
   /**
    * A converter for for assignments from a string value to a float value.
    */
-  public static class StringToFloatAssignmentConverter implements Converter<Map.Entry<String, Float>> {
+  public static class StringToFloatAssignmentConverter extends Converter.Contextless<Map.Entry<String, Float>> {
     private static final AssignmentConverter baseConverter = new AssignmentConverter();
 
     @Override

src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java

@@ -46,6 +46,7 @@
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.NoSuchElementException;
 import javax.annotation.Nullable;
 import org.apache.commons.pool2.PooledObject;
 import org.junit.Before;
@@ -114,21 +115,42 @@ private ResourceHandle acquire(double ram, double cpu, int tests)
     return acquire(ram, cpu, tests, ResourcePriority.LOCAL);
   }
 
-  private ResourceHandle acquire(double ram, double cpu, int tests, ImmutableMap<String, Float> extraResources, String mnemonic)
+  private ResourceHandle acquire(double ram, double cpu, int tests, String mnemonic)
       throws InterruptedException, IOException {
 
     return rm.acquireResources(
         resourceOwner,
-        ResourceSet.createWithWorkerKey(ram, cpu, tests, extraResources, createWorkerKey(mnemonic)),
+        ResourceSet.createWithWorkerKey(ram, cpu, tests, createWorkerKey(mnemonic)),
+        ResourcePriority.LOCAL);
+  }
+
+  private ResourceHandle acquire(double ram, double cpu, ImmutableMap<String, Float> extraResources, int tests, ResourcePriority priority)
+      throws InterruptedException, IOException, NoSuchElementException {
+    return rm.acquireResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests), priority);
+  }
+
+  private ResourceHandle acquire(double ram, double cpu, ImmutableMap<String, Float> extraResources, int tests)
+      throws InterruptedException, IOException, NoSuchElementException {
+    return acquire(ram, cpu, extraResources, tests, ResourcePriority.LOCAL);
+  }
+
+  private ResourceHandle acquire(double ram, double cpu, ImmutableMap<String, Float> extraResources, int tests, String mnemonic)
+      throws InterruptedException, IOException, NoSuchElementException {
+
+    return rm.acquireResources(
+        resourceOwner,
+        ResourceSet.createWithWorkerKey(ram, cpu, extraResources, tests, createWorkerKey(mnemonic)),
         ResourcePriority.LOCAL);
   }
 
   private void release(double ram, double cpu, int tests) throws IOException, InterruptedException {
-    rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests), /* worker=*/ null);
+    rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests), /* worker= */ null);
   }
 
-  private void release(double ram, double cpu, int tests, ImmutableMap<String, Float> extraResources) {
-    rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests));
+  private void release(double ram, double cpu, ImmutableMap<String, Float> extraResources, int tests)
+      throws InterruptedException, IOException {
+    rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests),
+        /* worker= */ null);
   }
 
   private void validate(int count) {
@@ -179,11 +201,11 @@ public void testOverBudgetRequests() throws Exception {
     release(0, 0, bigTests);
     assertThat(rm.inUse()).isFalse();
 
-    // Ditto, for extra resources (even if they don't exist in the available resource set):
-    ImmutableMap<String, Float> bigExtraResources = ImmutableMap.of("gpu", 10.0f, "fancyresource", 10.0f, "nonexisting", 10.0f);
-    acquire(0, 0, 0, bigExtraResources);
+    // Ditto, for extra resources:
+    ImmutableMap<String, Float> bigExtraResources = ImmutableMap.of("gpu", 10.0f, "fancyresource", 10.0f);
+    acquire(0, 0, bigExtraResources, 0);
     assertThat(rm.inUse()).isTrue();
-    release(0, 0, 0, bigExtraResources);
+    release(0, 0, bigExtraResources, 0);
     assertThat(rm.inUse()).isFalse();
   }
 
@@ -271,20 +293,21 @@ public void testThatExtraResourcesCannotBeOverallocated() throws Exception {
     assertThat(rm.inUse()).isFalse();
 
     // Given a partially acquired extra resources:
-    acquire(0, 0, 1, ImmutableMap.of("gpu", 1.0f));
+    acquire(0, 0, ImmutableMap.of("gpu", 1.0f), 1);
 
     // When a request for extra resources is made that would overallocate,
     // Then the request fails:
-    TestThread thread1 = new TestThread(() -> assertThat(acquireNonblocking(0, 0, 0, ImmutableMap.of("gpu", 1.1f))).isNull());
+    TestThread thread1 = new TestThread(() -> acquire(0, 0, ImmutableMap.of("gpu", 1.1f), 0));
     thread1.start();
-    thread1.joinAndAssertState(10000);
+    AssertionError e = assertThrows(AssertionError.class, () -> thread1.joinAndAssertState(1000));
+    assertThat(e).hasCauseThat().hasMessageThat().contains("is still alive");
   }
 
   @Test
   public void testHasResources() throws Exception {
     assertThat(rm.inUse()).isFalse();
     assertThat(rm.threadHasResources()).isFalse();
-    acquire(1, 0.1, 1, ImmutableMap.of("gpu", 1.0f));
+    acquire(1, 0.1, ImmutableMap.of("gpu", 1.0f), 1);
     assertThat(rm.threadHasResources()).isTrue();
 
     // We have resources in this thread - make sure other threads
@@ -305,15 +328,15 @@ public void testHasResources() throws Exception {
               assertThat(rm.threadHasResources()).isTrue();
               release(0, 0, 1);
               assertThat(rm.threadHasResources()).isFalse();
-              acquire(0, 0, 0, ImmutableMap.of("gpu", 1.0f));
+              acquire(0, 0, ImmutableMap.of("gpu", 1.0f), 0);
               assertThat(rm.threadHasResources()).isTrue();
-              release(0, 0, 0, ImmutableMap.of("gpu", 1.0f));
+              release(0, 0, ImmutableMap.of("gpu", 1.0f), 0);
               assertThat(rm.threadHasResources()).isFalse();
             });
     thread1.start();
     thread1.joinAndAssertState(10000);
 
-    release(1, 0.1, 1, ImmutableMap.of("gpu", 1.0f));
+    release(1, 0.1, ImmutableMap.of("gpu", 1.0f), 1);
     assertThat(rm.threadHasResources()).isFalse();
     assertThat(rm.inUse()).isFalse();
   }
@@ -704,6 +727,18 @@ private CyclicBarrier startAcquireReleaseThread(ResourcePriority priority) {
     return sync;
   }
 
+  @Test
+  public void testNonexistingResource() throws Exception {
+    // If we try to use nonexisting resource we should return an error
+    TestThread thread1 =
+        new TestThread(
+            () -> {
+              assertThrows(NoSuchElementException.class, () -> acquire(0, 0, ImmutableMap.of("nonexisting", 1.0f), 0));
+            });
+    thread1.start();
+    thread1.joinAndAssertState(1000);
+  }
+
   @Test
   public void testAcquireWithWorker_acquireAndRelease() throws Exception {
     int memory = 100;