fixup Implement extra resources for test actions

Addressed comments on bazelbuild#13996

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

Author: Drew Macrae

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/ResourceSet.java

@@ -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/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

@@ -407,6 +407,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

@@ -114,21 +114,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 {
+    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 {
+    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 {
+
+    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) {
@@ -181,9 +202,9 @@ public void testOverBudgetRequests() throws Exception {
 
     // 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);
+    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 +292,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(10000));
+    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 +327,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();
   }